[Fix] lm_head lora save by Datta0 · Pull Request #4106 · unslothai/unsloth

Datta0 · 2026-02-25T04:51:49Z

Fixes : #4098
lm_head should ideally be a module_to_save and not target_module
This is confirmed to work by @marcandrelarochelle the OP of the issue

Needs: unslothai/unsloth-zoo#515

gemini-code-assist · 2026-02-25T04:52:02Z

Summary of Changes

Hello @Datta0, I'm Gemini Code Assist¹! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

This pull request addresses an issue where the lm_head module was not being correctly handled during LoRA model saving. By explicitly defining lm_head as an accepted module and removing the previous workaround, the change ensures that lm_head is properly treated as a module to be saved, leading to more robust and correct LoRA model configurations.

Highlights

LM Head Module Handling: The lm_head module has been explicitly added to the accepted_modules set, ensuring it is correctly recognized for LoRA configurations.
Simplified LoRA Configuration Logic: Removed the conditional logic within get_peft_model that previously handled lm_head by moving it from target_modules to modules_to_save, streamlining the configuration process.

🧠 New Feature in Public Preview: You can now enable Memory to help Gemini Code Assist learn from your team's feedback. This makes future code reviews more consistent and personalized to your project's style. Click here to enable Memory in your admin console.

Activity

The fix has been confirmed to work by @marcandrelarochelle, the original reporter of the issue [Bug] lm_head is not trained using LoRA and merging is broken #4098.

Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point by creating a comment using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands on the current page.

Feature	Command	Description
Code Review	`/gemini review`	Performs a code review for the current pull request in its current state.
Pull Request Summary	`/gemini summary`	Provides a summary of the current pull request in its current state.
Comment	@gemini-code-assist	Responds in comments when explicitly tagged, both in pull request comments and review comments.
Help	`/gemini help`	Displays a list of available commands.

Customization

To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here.

You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.

Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

gemini-code-assist

Code Review

This pull request aims to fix the saving mechanism for lm_head when using LoRA. The changes involve adding lm_head to the accepted_modules list, allowing it to be directly specified as a target_module, and removing the previous logic that would automatically move lm_head from target_modules to modules_to_save. However, this appears to contradict the pull request description, which states that lm_head should ideally be a module_to_save and not a target_module. Clarification on the intended design for lm_head's role in PEFT configuration is needed to ensure consistency between the code's behavior and its documentation.

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.

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

chatgpt-codex-connector

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: abc9dbad4b

ℹ️ About Codex in GitHub

Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you

Open a pull request for review
Mark a draft as ready
Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".

chatgpt-codex-connector · 2026-03-03T14:04:00Z

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


Keep lm_head out of LoRA target_modules

By removing the lm_head special-case here, lm_head now goes into final_modules and is trained as a LoRA target instead of being moved to modules_to_save. That breaks merged exports: the merge code only applies _merge_lora to LLAMA_WEIGHTS and never merges LoRA deltas for lm_head (see unsloth/save.py where LLAMA_WEIGHTS excludes lm_head and lm_head.weight is copied directly), so merged_16bit/merged_4bit outputs silently drop lm_head training whenever users include lm_head in target_modules.

Useful? React with 👍 / 👎.

chatgpt-codex-connector

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: a07a417354

ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

Open a pull request for review
Mark a draft as ready
Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".

chatgpt-codex-connector · 2026-03-03T14:32:54Z

-                    modules_to_save.append("lm_head")
-
-            elif module == "embed_tokens":
+            if module == "embed_tokens":


Keep lm_head out of LoRA target modules

By no longer special-casing lm_head here, lm_head now falls through into final_modules and is trained as a LoRA target instead of being forced into modules_to_save. In merged export paths this silently drops lm_head training: unsloth/save.py only calls _merge_lora for names in LLAMA_WEIGHTS (unsloth/save.py:84-92, 641-644) and then writes lm_head.weight directly (690-692), so any LoRA delta on lm_head is not merged into merged_16bit/merged_4bit outputs.

Useful? React with 👍 / 👎.

* 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>

Fix lm_head lora save

a8eb93b

Datta0 requested review from danielhanchen and mmathew23 as code owners February 25, 2026 04:51

Datta0 mentioned this pull request Feb 25, 2026

[Fix] lm_head lora save unslothai/unsloth-zoo#515

Merged

gemini-code-assist Bot reviewed Feb 25, 2026

View reviewed changes

Comment thread unsloth/models/llama.py

Datta0 and others added 3 commits March 3, 2026 13:57

Merge branch 'main' into lora_save

557386b

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

abc9dba

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

chatgpt-codex-connector Bot reviewed Mar 3, 2026

View reviewed changes

Merge branch 'main' into lora_save

a07a417

danielhanchen merged commit af59aeb into unslothai:main Mar 3, 2026
1 check passed

chatgpt-codex-connector Bot reviewed Mar 3, 2026

View reviewed changes

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

[Fix] lm_head lora save#4106

[Fix] lm_head lora save#4106
danielhanchen merged 5 commits into
unslothai:mainfrom
Datta0:lora_save

Datta0 commented Feb 25, 2026 •

edited

Loading

Uh oh!

gemini-code-assist Bot commented Feb 25, 2026

Uh oh!

gemini-code-assist Bot left a comment

Uh oh!

Uh oh!

chatgpt-codex-connector Bot left a comment

Uh oh!

chatgpt-codex-connector Bot Mar 3, 2026

Uh oh!

Uh oh!

chatgpt-codex-connector Bot left a comment

Uh oh!

chatgpt-codex-connector Bot Mar 3, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Uh oh!

Conversation

Datta0 commented Feb 25, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

gemini-code-assist Bot commented Feb 25, 2026

Summary of Changes

Highlights

Footnotes

Uh oh!

gemini-code-assist Bot left a comment

Choose a reason for hiding this comment

Code Review

Uh oh!

Uh oh!

chatgpt-codex-connector Bot left a comment

Choose a reason for hiding this comment

💡 Codex Review

Uh oh!

chatgpt-codex-connector Bot Mar 3, 2026

Choose a reason for hiding this comment

Uh oh!

Uh oh!

chatgpt-codex-connector Bot left a comment

Choose a reason for hiding this comment

💡 Codex Review

Uh oh!

chatgpt-codex-connector Bot Mar 3, 2026

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Datta0 commented Feb 25, 2026 •

edited

Loading