From b7ec64c96fd056a4d9902b87cc9299e611e4407b Mon Sep 17 00:00:00 2001
From: Datta Nimmaturi <venkatadattasainimmaturi@gmail.com>
Date: Tue, 3 Mar 2026 20:00:13 +0530
Subject: [PATCH] [Fix] lm_head lora save (#4106)

* Fix lm_head lora save

* Fix _need_to_train_embeddings guard for lm_head LoRA targets

When lm_head is already in final_modules as a LoRA target, the
_need_to_train_embeddings block should not also add it to
modules_to_save. This prevents dual-wrapping (LoRA + modules_to_save
on the same module) which causes assertion failures downstream.

Check if embed_tokens/lm_head are already being trained as LoRA
targets before adding them to modules_to_save. Also prevents
duplicate entries with elif guards.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

---------

Co-authored-by: Daniel Han <danielhanchen@gmail.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
---
 unsloth/models/llama.py | 45 ++++++++++++++++++++---------------------
 1 file changed, 22 insertions(+), 23 deletions(-)

diff --git a/unsloth/models/llama.py b/unsloth/models/llama.py
index 6fe60cf940..6e2fdb57ca 100644
--- a/unsloth/models/llama.py
+++ b/unsloth/models/llama.py
@@ -2957,6 +2957,7 @@ class FastLlamaModel:
 
         accepted_modules = frozenset(
             (
+                "lm_head",
                 "q_proj",
                 "k_proj",
                 "v_proj",
@@ -2975,18 +2976,7 @@ class FastLlamaModel:
         train_embed_tokens = False
         final_modules = []
         for module in target_modules:
-            if module == "lm_head":
-                # logger.warning_once(
-                #     "Unsloth: `lm_head` should be placed in `modules_to_save` and not `target_modules`. "\
-                #     "Luckily, we shall do it for you!"
-                # )
-                train_lm_head = True
-                if modules_to_save is None:
-                    modules_to_save = ["lm_head"]
-                else:
-                    modules_to_save.append("lm_head")
-
-            elif module == "embed_tokens":
+            if module == "embed_tokens":
                 # logger.warning_once(
                 #     "Unsloth: `embed_tokens` should be placed in `modules_to_save` and not `target_modules`. "\
                 #     "Luckily, we shall do it for you!"
@@ -3011,23 +3001,32 @@ class FastLlamaModel:
 
         # Check if we added new tokens!
         if hasattr(model, "_need_to_train_embeddings"):
-            if not train_lm_head or not train_embed_tokens:
+            # Check if embed_tokens/lm_head are already being trained
+            # (either as LoRA targets in final_modules or via modules_to_save)
+            _embed_already_trained = (
+                train_embed_tokens or "embed_tokens" in final_modules
+            )
+            _lm_head_already_trained = train_lm_head or "lm_head" in final_modules
+            if not _lm_head_already_trained or not _embed_already_trained:
                 print(
                     "Unsloth: You added new tokens but did not specify if you wanted to "
                     "train the lm_head and embed_tokens.\nWe must turn it on for you."
                 )
-                train_lm_head = True
-                train_embed_tokens = True
 
-                if modules_to_save is None:
-                    modules_to_save = ["embed_tokens"]
-                else:
-                    modules_to_save.append("embed_tokens")
+                # Only add to modules_to_save if not already a LoRA target
+                if not _embed_already_trained:
+                    train_embed_tokens = True
+                    if modules_to_save is None:
+                        modules_to_save = ["embed_tokens"]
+                    elif "embed_tokens" not in modules_to_save:
+                        modules_to_save.append("embed_tokens")
 
-                if modules_to_save is None:
-                    modules_to_save = ["lm_head"]
-                else:
-                    modules_to_save.append("lm_head")
+                if not _lm_head_already_trained:
+                    train_lm_head = True
+                    if modules_to_save is None:
+                        modules_to_save = ["lm_head"]
+                    elif "lm_head" not in modules_to_save:
+                        modules_to_save.append("lm_head")
 
         # Check for Llama-3
         # if hasattr(model._saved_temp_tokenizer, "_using_llama3_template"):