Fix silent embedding corruption in add_new_tokens on padded-vocab mod (#923)

Author: BardiaKoopah

tests/test_add_new_tokens_padded_vocab.py

@@ -0,0 +1,134 @@
+# SPDX-License-Identifier: AGPL-3.0-only
+# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved.
+
+"""Regression test for #121: add_new_tokens must not shrink a padded embedding.
+
+Some models ship an embedding padded LARGER than the tokenizer vocab (Gemma3:
+262208 rows vs 262145 tokens). The old code resized to len(tokenizer), which
+SHRANK the matrix and silently destroyed the already-trained rows past the
+tokenizer length (no exception raised). The fix never resizes below the existing
+embedding, fills only the genuinely-new rows with the trained mean, keeps tied
+weights tied, and keeps config.vocab_size equal to the real matrix row count so
+the model still round-trips through save_pretrained/from_pretrained.
+
+Fully synthetic (tiny tied Llama, ~100-token tokenizer), CPU only, no download.
+"""
+
+import tempfile
+
+import pytest
+import torch
+
+transformers = pytest.importorskip("transformers")
+pytest.importorskip("tokenizers")
+
+from tokenizers import Tokenizer
+from tokenizers.models import WordLevel
+from tokenizers.pre_tokenizers import Whitespace
+from transformers import LlamaConfig, LlamaForCausalLM, PreTrainedTokenizerFast
+
+from unsloth_zoo.tokenizer_utils import add_new_tokens
+
+VOCAB_TOK = 100  # real tokens in the tokenizer
+PADDED = 128     # embedding rows shipped by the model (padding = [100, 128))
+
+
+def _build(vocab_tok, padded):
+    torch.manual_seed(0)
+    vocab = {f"tok{i}": i for i in range(vocab_tok)}
+    backend = Tokenizer(WordLevel(vocab=vocab, unk_token="tok0"))
+    backend.pre_tokenizer = Whitespace()
+    tokenizer = PreTrainedTokenizerFast(
+        tokenizer_object=backend, unk_token="tok0", pad_token="tok1",
+    )
+    cfg = LlamaConfig(
+        vocab_size=padded, hidden_size=16, intermediate_size=32,
+        num_hidden_layers=1, num_attention_heads=2, num_key_value_heads=2,
+        tie_word_embeddings=True,
+    )
+    model = LlamaForCausalLM(cfg)
+    model.eval()
+    return model, tokenizer
+
+
+def test_padded_embedding_preserves_trained_rows_and_places_new_tokens():
+    model, tokenizer = _build(VOCAB_TOK, PADDED)
+    emb = model.get_input_embeddings().weight
+
+    # Snapshot everything the fix must preserve, and mark the alignment padding
+    # rows [VOCAB_TOK, PADDED) with a distinctive sentinel so we can tell if a
+    # new token wrongly resolves to a leftover padding row.
+    trained_rows = emb.detach()[:VOCAB_TOK].clone()
+    with torch.no_grad():
+        emb[VOCAB_TOK:PADDED] = 7.0
+    sentinel_rows = emb.detach()[VOCAB_TOK:PADDED].clone()
+    assert emb.data_ptr() == model.get_output_embeddings().weight.data_ptr()
+
+    new_tokens = ["<newA>", "<newB>"]
+    add_new_tokens(model, tokenizer, new_tokens=new_tokens)
+
+    emb2 = model.get_input_embeddings().weight
+    head2 = model.get_output_embeddings().weight
+
+    # 1. No shrink: all trained rows survive byte-identical.
+    assert emb2.shape[0] == PADDED, "padded embedding must not shrink"
+    assert torch.equal(emb2[:VOCAB_TOK], trained_rows), "trained rows corrupted"
+
+    # 2. New tokens land at their real tokenizer IDs, each pointing at its OWN
+    #    mean-initialised row (not a leftover padding row).
+    new_len = len(tokenizer)
+    for offset, tok in enumerate(new_tokens):
+        tid = tokenizer.convert_tokens_to_ids(tok)
+        assert tid == VOCAB_TOK + offset
+        assert tokenizer(tok, add_special_tokens=False).input_ids == [tid]
+        row = emb2[tid]
+        assert not torch.allclose(row, torch.full_like(row, 7.0)), \
+            "new token resolved to a leftover padding row"
+        # Row is real (finite, non-zero), i.e. mean-initialised.
+        assert torch.isfinite(row).all() and row.abs().sum() > 0
+
+    # 3. mean-init touched ONLY the genuinely-new rows: the leftover padding
+    #    beyond the new tokens is exactly as shipped.
+    assert torch.equal(emb2[new_len:PADDED], sentinel_rows[new_len - VOCAB_TOK:]), \
+        "alignment padding rows were overwritten"
+
+    # 4. Tied weights stay tied, and share storage at a new-token row.
+    assert emb2.data_ptr() == head2.data_ptr(), "tie broken by resize"
+    assert torch.equal(emb2[VOCAB_TOK], head2[VOCAB_TOK])
+
+    # 5. config.vocab_size tracks the matrix, so the model round-trips.
+    assert model.config.vocab_size == PADDED
+    with tempfile.TemporaryDirectory() as d:
+        model.save_pretrained(d)
+        reloaded = LlamaForCausalLM.from_pretrained(d)
+    assert reloaded.get_input_embeddings().weight.shape[0] == PADDED
+    assert reloaded.config.vocab_size == PADDED
+
+    # 6. Forward pass over the new IDs works.
+    with torch.no_grad():
+        logits = model(torch.tensor([[VOCAB_TOK, VOCAB_TOK + 1, 3, 4]])).logits
+    assert logits.shape[-1] == PADDED
+
+
+def test_non_padded_embedding_still_grows_normally():
+    # embedding == tokenizer length: the ordinary case must be unaffected.
+    model, tokenizer = _build(VOCAB_TOK, VOCAB_TOK)
+    trained_rows = model.get_input_embeddings().weight.detach()[:VOCAB_TOK].clone()
+
+    add_new_tokens(model, tokenizer, new_tokens=["<newA>", "<newB>"])
+
+    emb = model.get_input_embeddings().weight
+    head = model.get_output_embeddings().weight
+    assert emb.shape[0] == VOCAB_TOK + 2, "non-padded model must grow to fit"
+    assert torch.equal(emb[:VOCAB_TOK], trained_rows), "trained rows corrupted"
+    assert model.config.vocab_size == VOCAB_TOK + 2
+    assert emb.data_ptr() == head.data_ptr()
+    for offset in range(2):
+        row = emb[VOCAB_TOK + offset]
+        assert torch.isfinite(row).all() and row.abs().sum() > 0
+
+
+if __name__ == "__main__":
+    test_padded_embedding_preserves_trained_rows_and_places_new_tokens()
+    test_non_padded_embedding_still_grows_normally()
+    print("ok")

unsloth_zoo/tokenizer_utils.py

@@ -138,24 +138,35 @@ def add_new_tokens(
     old_length = len(tokenizer)
     tokenizer.add_tokens(new_tokens)
     new_vocab_length = len(tokenizer)
+    # Some models ship an embedding that is padded LARGER than the tokenizer
+    # vocab (e.g. Gemma3: 262208 embedding rows vs 262145 tokens). Resizing to
+    # new_vocab_length would then SHRINK the matrix and silently destroy the
+    # already-trained rows past the tokenizer length. Never resize below the
+    # existing embedding: grow only when the new tokens overflow the padding.
+    # new_vocab_length already includes the freshly added tokens, so do NOT add
+    # their count again here.
+    resize_target = max(old_input_length, old_output_length, new_vocab_length)
     # Also resizes lm_head as well!
-    model.resize_token_embeddings(new_vocab_length)
+    model.resize_token_embeddings(resize_target)
 
     # If we use interpolation, we interpolate between the mean embeddings and
     # the Word2Vec sum of the other vectors
     embedding_matrix = model.get_input_embeddings ().weight
     lm_head_matrix   = model.get_output_embeddings().weight
 
-    # Confirm sizes are correct
-    if embedding_matrix.shape[0] != new_vocab_length:
+    # Confirm sizes are correct. Compare against resize_target (not
+    # new_vocab_length): for padded embeddings the matrices keep their larger
+    # row count. The old check compared against new_vocab_length, so on a padded
+    # model it either never fired (masking the silent shrink) or fired spuriously.
+    if embedding_matrix.shape[0] != resize_target:
         raise RuntimeError(
             "Unsloth: Embedding matrix size did not get resized properly. Please file a bug report!"
         )
-    if lm_head_matrix.shape[0]   != new_vocab_length:
+    if lm_head_matrix.shape[0]   != resize_target:
         raise RuntimeError(
             "Unsloth: LM Head matrix size did not get resized properly. Please file a bug report!"
         )
-    if model.config.vocab_size   != new_vocab_length:
+    if model.config.vocab_size   != resize_target:
         raise RuntimeError(
             "Unsloth: Model's config vocab_size did not get resized properly. Please file a bug report!"
         )
@@ -181,10 +192,13 @@ def add_new_tokens(
                 lm_head_matrix  [old_length+j] = mean_lm_head_token
         pass
     else:
-        # Now set the new tokens to the mean!
+        # Now set the new tokens to the mean! Only touch the genuinely-new rows
+        # [old_length:new_vocab_length]; on a padded embedding the rows past
+        # new_vocab_length are the model's original alignment padding and must be
+        # left as shipped (the old [old_length:] slice overwrote those too).
         with torch.no_grad():
-            embedding_matrix[old_length:] = mean_embedding
-            lm_head_matrix  [old_length:] = mean_lm_head
+            embedding_matrix[old_length:new_vocab_length] = mean_embedding
+            lm_head_matrix  [old_length:new_vocab_length] = mean_lm_head
     pass
 
     # We set a flag to say we need to train embeddings
@@ -195,15 +209,18 @@ def add_new_tokens(
     pass
     internal_model._need_to_train_embeddings = True
 
-    # Fix up all vocab sizes
+    # Fix up all vocab sizes. Use resize_target, not len(tokenizer): config
+    # vocab_size must equal the actual embedding / lm_head row count or the model
+    # cannot round-trip through save_pretrained/from_pretrained (state_dict shape
+    # mismatch). For a padded embedding these differ.
     current_model = model
     while hasattr(current_model, "model") and hasattr(current_model, "config"):
         if hasattr(current_model.config, "vocab_size"):
-            current_model.config.update({"vocab_size" : len(tokenizer)})
+            current_model.config.update({"vocab_size" : resize_target})
         current_model = current_model.model
     if hasattr(current_model, "model") and hasattr(current_model, "config"):
         if hasattr(current_model.config, "vocab_size"):
-            current_model.config.update({"vocab_size" : len(tokenizer)})
+            current_model.config.update({"vocab_size" : resize_target})
     pass
 
     # Must tie lm_head and embed_tokens if they are tied!