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

[romitjain]

Feature request

I wanted to describe an unexpected (or unusual) behaviour:

If we pass either of the embedding layers or the LM head in modules to save for a model whose embeddings are tied (tie_word_embeddings = True), peft does the following:

1. It (correctly) adds adaptors to mentioned layers (in case both layers are added, then separately), which essentially results in an untied layer (divergence from the original model architecture)
2. Trains them separately, and the resultant model has different weights for both the embedding layer and the LM head. While storing the state dict, it stores different weight parameters
3. But most critically, it does not update the model config to set tie_word_embeddings to be False. It leads to unexpected downstream behaviour in case the model is loaded and merged.
   3.1. If we load the model again via PeftModel.from_pretrained, it works as expected
   3.2. If we merge it and load again using AutoModel.from_pretrained, it ties the embedding layer and the LM head. This results in nonsensical generations.

There have been several reported issues and solutions in the past:

1. With tied embeddings adapter merged to tied layers #2018
2. Lm_head layer Problem in gemma2 : 2b-it #2244
3. How to finetune embeddings and LM head as a single layer when they are tied? #1750

One of the solves (that I am aware of)

1. Warn if using tied target module with tie_word_embeddings #2025 - But this just adds the warning

Your contribution

There are 2 possible solutions I see (applicable only for models with tie_word_embeddings = True

1. If either of embedding layer or the lm_head is added as a module to save, add an adaptor to both the modules to respect the original model's config and tie them explicitly. (or tie the adaptor)
2. Update the model config and mark tie_word_embeddings = False. This will solve the issue, but all LoRA finetunes will diverge in the model config from the base model

If we can agree on the right solution, I can raise a PR to solve this.

[BenjaminBossan]

Thanks for reporting this issue and giving a detailed description and thinking about possible solutions. I agree that how this is handled can be messy in PEFT.

First, let's establish a common basis:

If we pass either of the embedding layers or the LM head in modules to save for a model whose embeddings are tied (tie_word_embeddings = True), peft does the following:

1. It adds adaptors to both layers (in case both layers are added, then separately), which essentially results in an untied layer (divergence from the original model architecture)

Hmm, is this really what happens?

from transformers import AutoModelForCausalLM
from peft import LoraConfig, get_peft_model

model_id = "google/gemma-2-2b-it"
model = AutoModelForCausalLM.from_pretrained(model_id)
assert model.config.tie_word_embeddings

config = LoraConfig(modules_to_save=["embed_tokens"])
model = get_peft_model(model, config)


print(type(model.base_model.model.model.embed_tokens))
# prints <class 'peft.utils.other.ModulesToSaveWrapper'>
print(type(model.base_model.model.lm_head))
# prints <class 'torch.nn.modules.linear.Linear'>

So at least in this case, only one of the two becomes a ModulesToSaveWrapper.

There are 2 possible solutions I see (applicable only for models with tie_word_embeddings = True

1. If either of embedding layer or the lm_head is added as a target module, mark them both as target modules to respect the original model's config

Now you mention "target module", not "modules to save", so did you mean those? Even so, I cannot observe adapters being added to both layers (using LoraConfig(target_modules=["lm_head"])). So let's first work out a small example so that we can be sure we're talking about the same thing.

2. Update the model config and mark tie_word_embeddings = False. This will solve the issue but all LoRA finetunes will diverge in the model config from the base model

Ideally, I'd like to avoid this, as we want to preserve the original architecture of the base model.

[romitjain]

@BenjaminBossan Thanks for the reply.

I think I mixed a couple of things. Let me clarify -

Firstly,

So at least in this case, only one of the two becomes a ModulesToSaveWrapper.

Yes, this is expected. What I should have written is

"It (correctly) adds adaptors to mentioned layers (in case both layers are added, then separately), which essentially results in an untied layer (divergence from the original model architecture)".

Secondly,

Now you mention "target module", not "modules to save", so did you mean those? Even so, I cannot observe adapters being added to both layers (using LoraConfig(target_modules=["lm_head"])). So let's first work out a small example so that we can be sure we're talking about the same thing.

Sorry, I meant to say "modules to save". I haven't tried "target modules" yet.

---

The example that you shared actually captures the crux of the issue perfectly. Users may forget (or ignore the warning) to save both the modules in case the embeddings are tied and since the adaptor is added only to one of the layer (out of lm head and embedding), the merged model will have an unexpected behaviour.

I have updated my description as well.

[BenjaminBossan]

Thanks for clarifying.

1. If either of embedding layer or the lm_head is added as a module to save, add an adaptor to both the modules to respect the original model's config and tie them explicitly. (or tie the adaptor)

Just to be clear, you would want to add a modules_to_save for the layer with the tied weight automatically, and then tie the weight of the ModulesToSaveWrapper? I think this sounds like a reasonable approach as it preserves the intent of having the two layers being tied. However, I wonder if such a change would break existing code. It could be that users have added workarounds that would stop working. We should be careful here.

Maybe another suggestion would be to keep the existing behavior, while facilitating doing the "right thing" for the users. This could be via better documentation, improved warnings, and/or examples. Maybe we can also add something to the PEFT code to facilitate this, by e.g. overriding tie_weights to do the work you proposed in solution 1, without changing the current default behavior.

[romitjain]

@BenjaminBossan

Just to be clear, you would want to add a modules_to_save for the layer with the tied weight automatically, and then tie the weight of the ModulesToSaveWrapper?

Yes.

However, I wonder if such a change would break existing code. It could be that users have added workarounds that would stop working. We should be careful here.

Agreed, that this is possible. As a user, I had explicitly updated the model config so that all downstream things work as expected. Other than that, the user might explicitly tie the layers again. In both cases, I think we are helping the users. We can add a warning of the updated behaviour.

by e.g. overriding tie_weights to do the work you proposed in solution 1, without changing the current default behavior.

I did not understand this specific point. Is tie_weights a specific function or a flag? If yes, that also sounds reasonable to me. We provide a flag to the user and based on their convinience, they can choose to enable/disable this.

[BenjaminBossan]

I did not understand this specific point. Is tie_weights a specific function or a flag? If yes, that also sounds reasonable to me. We provide a flag to the user and based on their convinience, they can choose to enable/disable this.

I meant the model.tie_weights() method on transformers models. We could decide that if users target tied layers on a PEFT model and then call this method, we "harmonize" the modules the way you suggested in your first solution. That way, we would keep the current behavior for backwards compatibility but also offer an easy way to do what most users probably expect in this situation.

An argument against this suggestion could be that what would be happening under the hood is not exactly weight tying. We could thus think of another, PEFT-specific name for this method.

[romitjain]

As a user - I would be more inclined towards a PEFT-specific handling. Something along the lines of ensure_weight_tieing in LoRA config, which may choose to call this or some PEFT-specific function. This PEFT-specific function can either tie the underlying layer or the adaptor.

Apart from modules_to_save, we might also need to handle cases for target_modules. Same case (either of the layers or both are set in target_modules).

WDYT?

[BenjaminBossan]

As a user - I would be more inclined towards a PEFT-specific handling. Something along the lines of ensure_weight_tieing in LoRA config, which may choose to call this or some PEFT-specific function. This PEFT-specific function can either tie the underlying layer or the adaptor.

Yes, a PEFT-specific handling sounds better. I'm still undecided about the exact API. Generally, compared to a config argument, I'd tend towards a method call, something like peft_model.ensure_weight_tying() because

1. There is precedence with model.tie_weights() from transformers
2. It would later allow us to add more arguments if needed
3. It is a more explicit "opt-in" for the user (when loading a config, this setting could be easily missed)
4. I would like to avoid too many arguments for edge cases on the already bloated LoraConfig

However, there is a big argument in favor of a config argument, namely that it's a hands-off solution. This can come in handy when e.g. using trl's SFTTrainer, where the user passes the PEFT config and trl initializes the PEFT model for you. What's your opinion?

Apart from modules_to_save, we might also need to handle cases for target_modules. Same case (either of the layers or both are set in target_modules).

Yes, ideally we can have a solution that covers both. I'm not sure what it would take to implement this.

[romitjain]

1. In general, I am also against config arguments 😄. But I think, given the current landscape with other libraries having a huge list of arguments (eg, trainer args), it seems the most convenient to me. Config args have better discoverability by the user. SFTTrainer can decide to set defaults without too many code changes. We can decide to expose the function too.
2. For covering both modules_to_save and target_modules, I can get started with whatever you feel is easier to implement first and figure out the best pattern during that. How does that sound?

[BenjaminBossan]

Sounds good. I'd say let's get started and then we can still change course if we encounter problems. Don't hesitate to create an early draft PR. And again, thanks for you contribution.

[romitjain]

Thanks @BenjaminBossan

I'll get started on this this week.

[BenjaminBossan]

There will probably be some hard to cover edge cases (non-transformers models, target_modules being a string for regex matching, etc.). Don't hesitate to just raise an error in that case, the solution doesn't need to cover 100%.

[github-actions]

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

[romitjain]

This is being worked upon