Proposal: make trainable tokens more flexible to support LMHead tuning #2792
Conversation
|
Thanks for the suggestion! I think that this is already supported by passing |
|
No there's a tiny issue (which is what my PR addresses) - passing the Since the LMHead (at least for some models like Llama3.2) is |
There was a problem hiding this comment.
No there's a tiny issue (which is what my PR addresses) - passing the lm_head triggers an error when it tries to access the embedding dim
Right, makes sense. I think supporting nn.Linear makes sense. Some comments below.
Let's also add some tests for this case, for example in tests/test_trainable_tokens.py.
| bias = getattr(self.base_layer, "bias", None) | ||
| result = F.linear( | ||
| input=x, | ||
| weight=W, | ||
| bias=bias, |
There was a problem hiding this comment.
I think we should use self.get_base_layer() instead to be consistent with update_layer(). If you want, you can update the lines above for the F.embedding call and instance check as well.
| if embed_dim is None: | ||
| embed_dim = weight.shape[-1] |
There was a problem hiding this comment.
Can you elaborate the purpose of the last case?
| if isinstance(targets, str): | ||
| targets = [targets] | ||
|
|
||
| # If embeddings are untied, also include the output embedding (lm head) module name | ||
| try: | ||
| tied_cfg = model_config.get("tie_word_embeddings", False) | ||
| tied_keys = getattr(self.model, "_tied_weights_keys", None) | ||
| are_tied = bool(tied_cfg and tied_keys is not None) | ||
| except Exception: | ||
| are_tied = False | ||
|
|
||
| if not are_tied and hasattr(self.model, "get_output_embeddings"): | ||
| out_emb = self.model.get_output_embeddings() | ||
| if out_emb is not None: | ||
| for name, module in self.model.named_modules(): | ||
| if module is out_emb: | ||
| targets.append(name) | ||
| break | ||
|
|
||
| peft_config.target_modules = list(dict.fromkeys(targets)) |
There was a problem hiding this comment.
What's the idea behind targeting the output embedding automatically in case of untied weights? I don't see the benefit and I think this is also breaking backward compatibility with existing checkpoints.
|
Thanks for the feedback, I'll follow up by the end of the week/next week! |
2 participants
In the case of untied Embed/LM heads, it seems natural to let the Trainable Tokens implementation also support LM Heads (especially in the case of reserved special tokens). Currently, my understanding is that it only allows you to fine tune the input embeddings (unless I'm horribly mistaken?). Is there any reason to restrict it to just input embeddings? I hacked my local PEFT install to lift this restriction, and it could be broadly useful for others as well.
I didn't do any rigorous testing, so let me know if I'm missing anything obvious. I'd be happy to help out with this if there's interest.