support multiple LoRAs in batched inference scenario

[pacman100]

What does this PR do?

1. Support multiple LoRAs in batched inference setting. Here, we create sub-batches via logic groupby(adapter_name).

How to use:

[To Do]

ToDos:

• Tests
• Example Notebook
• Documentation
• Support all the layers types such as Conv, Embedding, Quantized 8-bit and 4-bit layers.

[HuggingFaceDocBuilderDev]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint.

[BenjaminBossan]

Thank you for tackling this Sourab. I think this is a feature that many users would find useful and having support for this being added should be high priority.

I'm just leaving some early comments for now, so that we can discuss the design:

I was also thinking about how to solve this problem. Especially, how to pass the information about which adapter to use for which sample to the forward method of the respective adapters. The solution presented here is that this information is part of the kwargs when calling the PeftModel/LoraModel, then it is immediately removed from the kwargs and saved as an attribute on the adapter layer, where it is retrieved during the forward call.

I'm not a big fan of this approach, as it works via side-effects. This makes it hard to reason about and more difficult debug. Furthermore, it requires a lot of care to handle correctly. As an example, say a user calls generate with adapter_indices, then they are set as an attribute. Let's assume that somewhere in the call, there is an error. The adapter_indices won't be cleaned up. Next time the user calls without adapter_indices, the attribute is still set, thus they get incorrect outputs without any indication that something went wrong. So at the very least, we have to take more care of doing the clean up correctly.

Another potential issue is that at the moment, adapter_indices is cleaned up at the end of the forward call of the adapter layer. But what happens if a single forward or generate call on the LoraModel/PeftModel requires multiple forward calls for each layer? Then only the first call is correct and the subsequent ones ignore the information.

Overall, this approach looks very brittle to me and if possible, I would like to find a better approach.

I assume that you also considered just passing down the kwargs with the adapter_indices but that this wouldn't work, because transformers/some model architectures cannot handle arbitrary arguments being added. So this more straightforward way is not possible.

An alternative idea could be to work with pre-forward hooks using with_kwargs=True, as they should allow us to add more arguments to the kwargs before they are passed to forward. This is still a side-effect with all the associated problems, and it would still require us to ensure that the hooks are properly cleaned up, but at least we get the handles, which only require a .remove() call to be cleaned up.

Other than that, I'm open for ideas and discussions, maybe there is a better solution we just haven't considered yet.

[review-notebook-app]

Check out this pull request on 

See visual diffs & provide feedback on Jupyter Notebooks.

---

Powered by ReviewNB

[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.

[BenjaminBossan]

No, bad bot!

[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.

[BenjaminBossan]

not stale...

[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.