How to finetune embeddings and LM head as a single layer when they are tied?

[GokulNC]

I am looking to LoRA-finetune models like Gemma, which have tied embeddings. But, I would also like to have the shared embeddings as trainable (the common embedding table corresponding to both input and output embeddings of the network).

How do I achieve this?

---

Note: Passing both ["embed_tokens","lm_head"] to modules_to_save will result in untying them, because PEFT will create separate tensor copies. Passing only ["embed_tokens"] will result in only the input embeddings trainable (by making a separate PEFT copy), while the output embeddings being as it is (the original tensor).

[BenjaminBossan]

One possibility that you could try is to not add the embeddings to modules_to_save but instead just LoRA-tune them by adding them to target_modules. This could be especially useful for Gemma models, since they have huge embedding layers, so fully fine-tuning them pushes the number of trainable parameters up by a lot.

Another possibility (untested) is to try to manually tie the weights after initializing the PEFT models. So something along the lines of:

config = LoraConfig(..., modules_to_save=["embed_tokens", "lm_head"])
model = get_peft_model(model, config)
# exact names depend on architecture:
model.base_model.model.model.decoder.embed_tokens.modules_to_save["default"].weight = model.base_model.model.lm_head.modules_to_save["default"].weight

[github-actions[bot]]

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.

[VarunGumma]

Hi @GokulNC, do you have any leads on this? We are interested to try this for IT2-Dist models as well, where the embedding and output-projection (lm_head) are tied in the decoder.