ENH: LoRA support for dynamically dispatching to custom layers

[BenjaminBossan]

Resolves #1867

Description

This is an experimental feature with a private API for now. If this feature finds adoption, I will work on adding an official API.

With this PR, we allow users to register their own LoRA layer types. This way, they can add their own support for hitherto unsupported layer types, say nn.Conv3d or nn.LSTM. Without this PR, they can only do that by creating a PR on PEFT with support for this new type and getting it merged.

The custom dispatch mechanism also allows users to override existing layer type mapping. This way, they can, for instance, provide their own lora.Linear layer type, instead of using the one from PEFT, to adapt nn.Linear layers.

Implementation

The implementation required only very few changes because we already have a mechanism for dynamic dispatching for LoRA. It is currently used, for instance, to dynamically add quantized target layers in case the right quantization library is installed.

This existing mechanism is now extended to include user provided LoRA layers if those were passed. These are checked first before checking the default PEFT supported layers.

What's missing for this to become an official API?

Right now, the main reason why this cannot be an official API is the question of how to persist the config. In the current implementation, we add an attribute that is a mapping from target layer type to LoRA layer type:

config._custom_modules == {CustomBaseLayer: CustomLoraLayer}

The entries of this dict are Python classes. Therefore, they cannot be json-serialized. We could think of possible solutions how to serialize and deserialize custom Python objects, but this is not trivial and potentially a security risk. Thus I would only really start working on this if the demand is sufficiently high. At that point, I would also add a public API instead of requiring the use of a private API.

As is, users can still save and load PEFT models with custom LoRA layers, they only need to add two lines of code to their scripts, as documented.

We could also think about adding support for methods other than LoRA. However, this would require to implement the dynamic dispatch mechanism for those other methods, which right now only exists for LoRA.

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

[younesbelkada]

Thanks a lot for this clean integration, docs and API ! Left only one nit but overall looks great !

[younesbelkada]

here we should IMO state that the signature of the init method should have base_layer and adapter_name in the correct order otherwise the API will fail

[BenjaminBossan]

Good point, I added an entry for the __init__.

[stevhliu]

Very cool! Left a comment about improving the introductory section a bit 😄