From 27d2d1df495928fbdaf2a3b356984c641b79aa7c Mon Sep 17 00:00:00 2001
From: Daniel Han <danielhanchen@gmail.com>
Date: Thu, 31 Oct 2024 16:25:26 -0700
Subject: [PATCH] Tied weights

---
 unsloth/models/_utils.py | 30 ++++++++++++++++++++++++++++--
 unsloth/models/llama.py  |  7 ++++++-
 2 files changed, 34 insertions(+), 3 deletions(-)

diff --git a/unsloth/models/_utils.py b/unsloth/models/_utils.py
index 0539e255ea..e1ed649560 100644
--- a/unsloth/models/_utils.py
+++ b/unsloth/models/_utils.py
@@ -479,6 +479,32 @@ def patch_tokenizer(model, tokenizer):
     if model is not None:
         model.config.update({"unsloth_version" : __version__})
 
+    # First remove pad and unk tokens if they are known to be BOS / EOS
+    possible_bad_tokens = (
+        "<|endoftext|>",
+        "<|im_start|>",
+        "<|im_end|>",
+        "<|begin_of_text|>",
+        "<|end_of_text|>",
+        "<s>",
+        "</s>",
+    )
+    input_ids = tokenizer(list(possible_bad_tokens), add_special_tokens = False).input_ids
+    possible_bad_tokens = frozenset(token for token, input_id in zip(possible_bad_tokens, input_ids) if len(input_id) == 1)
+
+    if hasattr(tokenizer, "pad_token") and tokenizer.pad_token in possible_bad_tokens:
+        print(f"Unsloth: Pad token was {tokenizer.pad_token} which is not a good idea. We shall fix this.")
+        tokenizer.pad_token = None
+    pass
+
+    has_bad_unk_token = False
+    if hasattr(tokenizer, "unk_token") and tokenizer.unk_token in possible_bad_tokens:
+        print(f"Unsloth: Unk token was {tokenizer.unk_token} which is not a good idea. We shall fix this.")
+        tokenizer.unk_token = None
+        has_bad_unk_token = True
+    pass
+
+    # Now check pad token again
     bad_pad_token = False
     if hasattr(tokenizer, "pad_token") and tokenizer.pad_token is not None:
         # Check if pad_token is not the same as eos_token otherwise the loss will ignore it!!
@@ -492,13 +518,13 @@ def patch_tokenizer(model, tokenizer):
     # Check if unknown token is broken
     fixed_unk_token = False
 
-    if hasattr(tokenizer, "unk_token") and tokenizer.unk_token is not None:
+    if (hasattr(tokenizer, "unk_token") and tokenizer.unk_token is not None) or has_bad_unk_token:
 
         eos_token = getattr(tokenizer, "eos_token", None)
         bos_token = getattr(tokenizer, "bos_token", None)
 
         old_unk_token = tokenizer.unk_token
-        if old_unk_token == eos_token or old_unk_token == bos_token:
+        if (old_unk_token == eos_token) or (old_unk_token == bos_token) or has_bad_unk_token:
             has_broken_unk = True
             # Use the unicode replacement characters
             possible_replacements = [
diff --git a/unsloth/models/llama.py b/unsloth/models/llama.py
index c0175bbfaf..8c4b7fd253 100644
--- a/unsloth/models/llama.py
+++ b/unsloth/models/llama.py
@@ -1838,7 +1838,8 @@ class FastLlamaModel:
         except: old_output_embedding = torch.zeros(0)
 
         # Check for tied weights as well
-        is_tied = old_input_embedding.data_ptr() == old_output_embedding.data_ptr()
+        is_tied = (old_input_embedding.data_ptr() == old_output_embedding.data_ptr()) \
+            or (model.config.tie_word_embeddings)
 
         # Check pad token's id -> we need to expand the embedding
         if len(tokenizer) > old_input_embedding.shape[0]:
@@ -1887,6 +1888,10 @@ class FastLlamaModel:
         else:
             correct_dtype = old_input_embedding.dtype
         pass
+
+        # Must tie lm_head and embed_tokens if they are tied!
+        # Otherwise error will occur on saving models ie use save_model
+        if is_tied: model.tie_weights()
         
         # Also patch all dtypes - BnB seems to not allocate the correct type?
         # BnB default dtype seems to be float16!