ENH Allow multiple adapters when using target_parameters

[BenjaminBossan]

Resolves #3340

Context

So far, we did not allow adding multiple LoRA adapters with target_parameters on the same layer. This was a known limitation. I have already attempted to solve this once (see #2710) but didn't have time to come up with a nice solution. As it was unclear if there was any real world need to support this, there was no further work on it since then. Now we know that there are practical application that may need it, so I resumed the work.

What doesn't work

The previous solution attempted to solve the issue by nesting the lora.ParamWrappers. So for adapters 'default' and 'other', we would end up having something like:

param_wrapper_default(param_wrapper_other(base_layer))

This was problematic. Not only could this result in very deep nesting, which is inefficient. What's worse is that state_dict key for 'other' would contain 'base_layer.' as an infix. Therefore, if we wanted to load the 'other' adapter without first loading the 'default' adapter, we would get a key mismatch.

We could also not simply strip out 'base_layer.' infix because we use nesting to deal with multiple nn.Parameters on the same module; to account for that, we need to keep the infixes.

Solution

The solution is pretty straightforward: We use the existing mechanism to store the parameters for the other adapter in the nn.ModuleDict. For this, we detect if the layer is already a ParamWrapper when adding the second adapter and update that layer instead of nesting it.

Caveat

Generally with MoE modules in Transformes, we have multiple nn.Parameters on the same module, e.g. gate_up_proj and down_proj. Therefore, we often want to target multiple parameters on the same target module. For LoRA, we usually only have one targeted parameter per LoRA layer. To allow targeting multiple nn.Parameters, we thus create multple LoRA layers, one per nn.Parameter, and nest them. The simple approach from this PR supports this as long as the different adapters target exactly the same nn.Parameters.

This solution does, however, not work with multiple adapters that target a different sets of parameters. This is because the information which parameter is targeted is not stored in the state_dict itself -- remember: When there is more than 1 parameter per layer, we solve this via nesting. So we rely on the nesting level 0 corresponding to parameter 0, nesting 1 to parameter 1, etc. Therefore, if we had different adapters targeting different parameters, we would lose that mapping 1:1 and would not be able to tell which parameter is meant to be targeted when loading a checkpoint. (I can't really remember why I didn't try the solution from this PR back then, but it could be for this exact limitation.)

If we wanted to allow targeting different parameters for each adapter, we would need larger changes. First, we would need to store the targeted parameter name in a dict, separately for each adapter name. Then, during serialization, we would need to mangle the parameter name into the state_dict keys to be able to restore the correct mapping. Moreover, since this changes the structure of the state_dict, we would need special logic to still support existing state_dicts that don't follow that logic. Overall, this would be a big change with unclear benefit. Thus, once again, I opt to postpone a full solution.

[BenjaminBossan]

Note: This check for the compatibility of the new config should have arguably been inside of _check_new_adapter_config. Therefore, the new checking code uses _check_new_adapter_config instead of running inside of _create_and_replace.

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

[githubnemo]

LGTM!