Add tokenizer checking + TinyLlama

README.md

@@ -12,6 +12,7 @@
 | [Colab A100 example](https://colab.research.google.com/drive/1YIPY_18xm-K0iJDgvNkRoJsgkPMPAO3G?usp=sharing) | [Colab A100 example](https://colab.research.google.com/drive/1SKrKGV-BZoU4kv5q3g0jtE_OhRgPtrrQ?usp=sharing) | (59 more examples if you scroll down) | [Kaggle Slim Orca example](https://www.kaggle.com/danielhanchen/unsloth-slimorca-t4-ddp) |
 
 * **NEW!** [DPO](https://arxiv.org/abs/2305.18290) support. [Free DPO Colab example](https://colab.research.google.com/drive/15vttTpzzVXv_tJwEk-hIcQ0S9FcEWvwP?usp=sharing). [More info](#DPO).
+* **NEW!** [TinyLlama](https://github.com/jzhang38/TinyLlama) on 3T tokens. [Free Colab example](https://colab.research.google.com/drive/1AZghoNBQaMDgWJpi4RbffGM1h6raLUj9?usp=sharing). We also show automatic RoPE Scaling extending TinyLlama from 2048 to 4096 tokens!
 * Supports Llama, Yi, Mistral, CodeLlama, Qwen (llamafied), Deepseek and their derived models (Open Hermes etc).
 * All kernels written in [OpenAI's Triton](https://openai.com/research/triton) language. **Manual backprop engine**.
 * **0% loss in accuracy** - no approximation methods - all exact.
@@ -95,10 +96,11 @@ fourbit_models = [
     "unsloth/llama-2-7b-bnb-4bit",
     "unsloth/llama-2-13b-bnb-4bit",
     "unsloth/codellama-34b-bnb-4bit",
+    "unsloth/tinyllama-bnb-4bit",
 ]
 # Load Llama model
 model, tokenizer = FastLanguageModel.from_pretrained(
-    model_name = "unsloth/mistral-7b", # Supports Llama, Mistral - replace this!
+    model_name = "unsloth/mistral-7b-bnb-4bit", # Supports Llama, Mistral - replace this!
     max_seq_length = max_seq_length,
     dtype = None,
     load_in_4bit = True,

unsloth/models/_utils.py

@@ -44,6 +44,7 @@ from xformers import __version__ as xformers_version
 __all__ = [
     "prepare_model_for_kbit_training",
     "patch_tokenizer",
+    "check_tokenizer",
     "xformers",
     "xformers_attention",
     "xformers_version",
@@ -112,3 +113,26 @@ def patch_tokenizer(model, tokenizer):
     pass
     return model, tokenizer
 pass
+
+
+def check_tokenizer(model, tokenizer):
+    # Checks tokenizer for out of bounds ids.
+    # Mainly a fix for https://huggingface.co/berkeley-nest/Starling-LM-7B-alpha
+    # where <sep> had token id=32002.
+    # See https://huggingface.co/berkeley-nest/Starling-LM-7B-alpha/discussions/25
+    special_tokens_map = tokenizer.special_tokens_map
+    max_embedding_size = model.model.embed_tokens.weight.shape[0]
+
+    for token_name, token_content in special_tokens_map.items():
+        if type(token_content) is not str: continue
+        token_ids = tokenizer([token_content], add_special_tokens = False, return_attention_mask = False)
+        token_ids = token_ids.input_ids[0][0]
+        if token_ids < 0 or token_ids >= max_embedding_size:
+            raise RuntimeError(
+                f"Unsloth: Extra special token `{token_content}` with id={token_ids} exceeds "\
+                f"the maximum vocabulary size of {max_embedding_size}. You must fix the tokenizer "\
+                "or else out of bounds memory accesses will occur."
+            )
+        pass
+    pass
+pass

unsloth/models/llama.py

@@ -791,6 +791,9 @@ class FastLlamaModel:
             assert(module in accepted_modules)
         pass
 
+        # We check the tokenizer first for errors
+        check_tokenizer(model, tokenizer)
+
         # Get LoRA
         lora_config = LoraConfig(
             r              = r,