Embeddings LoRA & TP

[michaelbenayoun]

To be merged after #3079 .

It adds support for LoRA embedding with TP.

cc @3outeille

[HuggingFaceDocBuilderDev]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

[BenjaminBossan]

Thanks for extending TP support to LoRA embedding layers. This PR mostly looks good as is, just a few comments, please check. Since they are not included in the PR CI, I ran the updated tests on my machine and they passed.

[BenjaminBossan]

Ah damn. Do you think the usage of SimpleNamespace is robust? Otherwise, we could consider using nn.utils.parametrize.

[michaelbenayoun]

Parametrize will not do much here, what we need is a module holding the weights so that the input and output functions can use it.

I changed SimpleNamespace, which is very basic and simple, with something a bit more elaborate. We create a nn.Module, holding the lora_embedding_A Parameter dict instead of individual weights. This way, even if the parameter changes, we still get the proper weight for a given adapter.

I think it is cleaner than the first approach.

[BenjaminBossan]

I changed SimpleNamespace, which is very basic and simple, with something a bit more elaborate. We create a nn.Module, holding the lora_embedding_A Parameter dict instead of individual weights.

That sounds better, but I think you didn't push that change yet.

Also, instead of monkey-patching lora_module._embed here, I think it's better to go into peft.tuners.lora.Embedding._embed and add branching code in there to check for tp_plan and call the corresponding helper functions there.

[michaelbenayoun]

I just pushed.
About your second suggestion, great idea. I just pushed it as well.

[BenjaminBossan]

Thanks for the refactor. This looks like the more robust solution to me. PR looks good now.