untrained tokens llama 3.1 base (#929)

unsloth/chat_templates.py

@@ -876,7 +876,7 @@ def get_chat_template(
 
     # Careful on Gemma
     # bos_token is a must or else losses become too high
-    if IS_GEMMA and not chat_template.startswith("{{ bos_token }}"):
+    if IS_GEMMA and not chat_template.startswith(("{{ bos_token }}", "{{- bos_token }}")):
         chat_template = "{{ bos_token }}" + chat_template
     pass
 
@@ -1553,7 +1553,7 @@ extra_eos_tokens = None,
 
     # Check jinja tempate for bos
     if always_bos_token:
-        if not jinja_template.startswith("{{ bos_token }}"):
+        if not jinja_template.startswith(("{{ bos_token }}", "{{- bos_token }}")):
             jinja_template = "{{ bos_token }}" + jinja_template
     pass
 

unsloth/tokenizer_utils.py

@@ -827,7 +827,27 @@ def fix_untrained_tokens(model, tokenizer, train_dataset, eps = 1e-16):
     # Get untrained tokens
     indicator_untrained1 = torch.amax(embedding_matrix, axis = 1) <= eps
     # Check lm_head as well
+
+    # Does NOT work for Llama 3.1!!
     indicator_untrained2 = torch.amax(lm_head_matrix,   axis = 1) <= eps
+
+    # We instead check for repeated vectors
+    lm_head_where = torch.where(indicator_untrained1)[0]
+    lm_head_bad = lm_head_matrix[lm_head_where]
+    lm_head_bad = lm_head_bad.cpu().numpy().round(3)
+    from collections import Counter
+    counter = Counter()
+    for row in lm_head_bad: counter[hash(row.data.tobytes())] += 1
+    counter = Counter({k: c for k, c in counter.items() if c >= 2})
+
+    lm_head_where = lm_head_where.cpu().numpy()
+    final_bad_lm_head = []
+    for j, row in enumerate(lm_head_bad):
+        if hash(row.data.tobytes()) in counter:
+            final_bad_lm_head.append(lm_head_where[j])
+    indicator_untrained2 = indicator_untrained2 | torch.zeros_like(indicator_untrained2)
+    indicator_untrained2[final_bad_lm_head] = True
+
     # Combine both checks
     indicator_untrained = indicator_untrained1 & indicator_untrained2