ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

[Fix] lm_head lora save (#4106)

Browse source 
* 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>
This commit is contained in:
Datta Nimmaturi 2026-03-03 20:00:13 +05:30 • committed by GitHub
commit b7ec64c96f
1 changed files with 22 additions and 23 deletions
Download patch file Download diff file Expand all files Collapse all files

45
unsloth/models/llama.py
View file

@ -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"):