Skip to content

[Bug] lm_head is not trained using LoRA and merging is broken #4098

@marcandrelarochelle

Description

@marcandrelarochelle

Latest version of Unsloth (Unsloth 2026.2.1)

After a full training run, I noticed the adapter_config.json, the key target_modules didn't include the lm_head, but the modules_to_save did include the lm_head, this resulted in this error when attempting to merge:

RuntimeError: Unsloth: Extracted keys = {'lm_head.weight'} do not match!

in the saving_utils from unsloth_zoo here: Exact line asserting the issue saving_utils.py#L303

Cause: the lm_head is silently getting filtered out or included somewhere within the code and causing issues down the line later on.

How to reproduce (Reproducible in Colab):

  • Load the peft module with the lm_head in target_modules (inside the get_peft_model)
    • if there's weight tying involved, it will automatically do the right thing (LoRA on top or not during training)

(Colab Example) Qwen3-4B-Thinking Notebook; modify the code for the get_peft_model

model = FastLanguageModel.get_peft_model(
    model,
    r = 32, # Choose any number > 0 ! Suggested 8, 16, 32, 64, 128
    target_modules = ["lm_head", "q_proj", "k_proj", "v_proj", "o_proj",
                      "gate_proj", "up_proj", "down_proj",],                 # Added "lm_head" to target_modules
    lora_alpha = 32,
    lora_dropout = 0, # Supports any, but = 0 is optimized
    bias = "none",    # Supports any, but = "none" is optimized
    # [NEW] "unsloth" uses 30% less VRAM, fits 2x larger batch sizes!
    use_gradient_checkpointing = "unsloth", # True or "unsloth" for very long context
    random_state = 3407,
    use_rslora = False,  # We support rank stabilized LoRA
    loftq_config = None, # And LoftQ
)
  • Train the model
  • Verify the adapter_config.json of the resulting trained LoRA adapter (to validate the lm_head shows, you'll see no lm_head in the target_modules within the adapter_config.json, but lm_head will be in the modules_to_save)
  • Attempt merging using save_pretrained_merged

Within the same notebook, at the Saving float16 for VLLM, you switch the first if clause to True

if True:
    model.save_pretrained_merged("qwen_finetune_16bit", tokenizer, save_method = "merged_16bit",)
  • and then the specified exception appears

RuntimeError: Unsloth: Extracted keys = {'lm_head.weight'} do not match!

Additional related information (What are the expected behaviors):

Metadata

Metadata

Assignees

No one assigned

    Labels

    Type

    No type
    No fields configured for issues without a type.

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions