Skip to content

Conversation

@www-spam
Copy link

Summary

This PR fixes the issue where merge_and_unload() produces broken models when adapters are applied to both embed_tokens and lm_head on models with tie_word_embeddings=True.

Resolves #2777

Problem

When a base model has tie_word_embeddings=True (e.g., Gemma, Llama):

  1. embed_tokens and lm_head share the same weight tensor
  2. Adapters can be applied to both layers (via modules_to_save or target_modules)
  3. After training, each layer has different adapter deltas
  4. merge_and_unload() merges both layers with their respective deltas
  5. Bug: The config still has tie_word_embeddings=True
  6. When the merged model is loaded with AutoModel.from_pretrained(), the lm_head weights are overwritten with embed_tokens weights due to weight tying
  7. Result: The merged lm_head weights are lost, causing degraded or garbage output

Solution

This PR modifies _unload_and_optionally_merge() in BaseTuner to:

  1. Detect if both embedding-like and lm_head-like modules have adapters
  2. Untie the weights by cloning lm_head.weight before merge
  3. Update config.tie_word_embeddings = False in all relevant config locations

This ensures that:

  • Merged weights are preserved for both layers
  • The saved model can be loaded correctly
  • Backward compatibility is maintained (no change when embeddings aren't both targeted)

Changes

  • src/peft/tuners/tuners_utils.py:

    • Added _untie_embedding_weights() helper method
    • Added _update_tie_word_embeddings_config() helper method
    • Added _has_adapters_on_both_embeddings() helper method
    • Modified _unload_and_optionally_merge() to auto-handle tied embeddings
  • tests/test_tie_word_embeddings_merge.py:

    • Added tests for tie_word_embeddings merge behavior

Test Plan

  • Tested with Gemma 3 4B model (tie_word_embeddings=True)
  • Verified merged model produces coherent output
  • Verified config.tie_word_embeddings is correctly set to False
  • Verified embed_tokens and lm_head have independent weights after merge
  • Unit tests added

Example

Before this fix:

# Model with tie_word_embeddings=True + adapters on embed_tokens and lm_head
merged = peft_model.merge_and_unload()
merged.save_pretrained("merged_model")

# Loading produces broken model
loaded = AutoModel.from_pretrained("merged_model")  # lm_head weights lost!

After this fix:

# Same setup
merged = peft_model.merge_and_unload()  # Auto-unties and updates config
merged.save_pretrained("merged_model")

# Loading works correctly
loaded = AutoModel.from_pretrained("merged_model")  # Works as expected

…pted

When LoRA is applied to both embed_tokens and lm_head on models with
tie_word_embeddings=True, merge_and_unload() now automatically:
- Detects if both layers have adapters
- Unties the weights before merging
- Sets config.tie_word_embeddings=False

This prevents the merged lm_head weights from being lost when the
model is reloaded.

Resolves huggingface#2777
@www-spam www-spam closed this Dec 31, 2025
@www-spam www-spam reopened this Dec 31, 2025
- Check both target_modules and modules_to_save when detecting adapters
  on embed_tokens and lm_head
- Always update config when adapters are on both layers (ModulesToSaveWrapper
  already unties weights, so we just need to update config)
- Update warning message for clarity
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Update the model config to tie_word_embeddings = False in case of lm_head or embedding layer update

1 participant