From 820cd4efef5d173f2200ef184b3d8856ba09d75b Mon Sep 17 00:00:00 2001
From: Daniel Han <danielhanchen@gmail.com>
Date: Thu, 31 Oct 2024 12:36:21 -0700
Subject: [PATCH] Tied weights

---
 unsloth/models/llama.py | 6 +++++-
 1 file changed, 5 insertions(+), 1 deletion(-)

diff --git a/unsloth/models/llama.py b/unsloth/models/llama.py
index c0175bbfaf..0e9b70a8b2 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()
+        is_tied = 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,9 @@ class FastLlamaModel:
         else:
             correct_dtype = old_input_embedding.dtype
         pass
+
+        # Finally tie them if needed!
+        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!