[WIP] Add LoRA multihead attention module

[BenjaminBossan]

First stab at adding LoRA support for nn.MultiheadAttention. See #761.

Todos:

• For now, only works with _qkv_same_embed_dim=True -- make it work with False too. _qkv_same_embed_dim=False is out of scope for this PR and can be added in a later PR if needed.
• Show that it works in a real world test: See user feedback on the issue.
• Unit tests
• Docs Apart from docstrings, I don't think anything else needs to be added

Update: I now also included the out_proj to apply LoRA to.

This is a simple test that I ran successfully with the PR in its current state:

import open_clip
import requests
import torch
from torch import nn
from peft import LoraConfig, get_peft_model
from PIL import Image
from peft.tuners.lora.layer import MultiheadAttention as PeftMha

model, preprocess = open_clip.create_model_from_pretrained('hf-hub:laion/CLIP-ViT-g-14-laion2B-s12B-b42K')
tokenizer = open_clip.get_tokenizer('hf-hub:laion/CLIP-ViT-g-14-laion2B-s12B-b42K')
peft_model = get_peft_model(model, config)
opt = torch.optim.SGD(peft_model.parameters(), 0.1)
print(len([m for m in peft_model.modules() if isinstance(m, PeftMha)]))  # 64 PEFT MHA layers
peft_model.print_trainable_parameters()  # trainable params: 2,588,672 || all params: 1,055,873,793 || trainable%: 0.24516869508096598

# text encoder
text = tokenizer(["a diagram", "a dog", "a cat"])
text_features = peft_model.encode_text(text)
loss = text_features.sum()
loss.backward()
opt.step()

# image encoder
url = 'http://images.cocodataset.org/val2017/000000039769.jpg'
image = Image.open(requests.get(url, stream=True).raw)
image = preprocess(image).unsqueeze(0)
image_features = model.encode_image(image)
image_features.sum().backward()
opt.step()

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

Nice work ! I left few preliminary comments, I think we can go for the _restore_weights approach for now as I don't see any other alternative

[younesbelkada]

Suggested change

is_target_conv_1d_layer: bool = False,

I don't think this is used?

[younesbelkada]

Suggested change

self.is_target_conv_1d_layer = is_target_conv_1d_layer

We can also just hard-code it to False

[younesbelkada]

do we need also to over-write the modules() method?

[BenjaminBossan]

Not needed, as modules calls named_modules under the hood. I added a comment to that effect.

[younesbelkada]

Why this has been removed?

[BenjaminBossan]

Sorry, forgot to put this into the description of the PR.

These lines are obsolete for some time now. They only apply when we unload the model (otherwise, the if does not match). Remember when we made the base_layer switch, we ensured that when unloading, we simply return the base_layer, no more need to create a new layer (say, a new nn.Linear when using lora.Linear) and replace the new layer's weight by the parent layer's weight. The base_layer already has the original weight. Therefore, these lines are unnecessary.

I removed them now because they were annoying with MultiheadAttention, because that layer has no weight attribute, so this line would fail.

[pacman100]

Thank you Benjamin for adding support for torch MHA layer in LoRA, interesting way to use merge, forward and unmerge logic!

[BenjaminBossan]

@younesbelkada Could I address all your concerns?

I pinged the user who wanted to test it on their case. When it comes to docs, I didn't really find a place where we list all supported layers, so no update needed really.

[BenjaminBossan]

Note: The test test_merge_layers for MHA fails. This is most likely because of an existing bug in how merging is implemented, see PR #1355. Once that is merged, the test should pass.

[ambroser53]

Just want to bump a bunch of the issues I've mentioned in #761 but specifically the problem with requires_grad reproducable in this repo

[bghira]

just wanted to bump this one because it's really the only way for tuning CLIP models after they are released.

[BenjaminBossan]

@bghira Do you happen to have a use case where you could test if this PR works and is working well enough speed-wise? I think the implementation could be ready to be merged but ideally we'd have someone with a real use case give it a try.

[bghira]

i do and i may be able to test it. stupid question but is the code example above complete? i dont see the hinge loss function

[BenjaminBossan]

stupid question but is the code example above complete? i dont see the hinge loss function

You mean the code right at the top? No, it's not complete at all, just a quick test to show that MHA is applied and the backward pass does not fail. This is not proper nor complete training code.

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

[mashijie1028]

@BenjaminBossan
Hi! I found that LoRA does not work for in_proj_weight in attn of open_clip. I was wondering how to fix this.
To be more specific, when I implement LoRA as follows:

lora_config = LoraConfig(
    r=16,
    target_modules=["in_proj_weight"],
    lora_alpha=32,
    lora_dropout=0.05
)

An error occurs as ValueError: Target modules {'in_proj_weight'} not found in the base model. Please check the target modules and try again.
But when I implement for out_proj, LoRA works fine!
Could you please tell me how to set target_modules in LoraConfig to implement LoRA on attn layers? Thanks!

By the way, I download peft as you mentioned before:

python -m pip install git+https://github.com/BenjaminBossan/peft.git@feat-add-lora-multihead-attention

(I report the same issue here)

[bghira]

not stale

[coding-kuku]

@BenjaminBossan
Hello,

I tried using the code from this PR, but I found that the batch_first attribute of nn.MultiheadAttention is lost after converting the model to a LoRA model. Below is the minimal code to reproduce the issue:

import torch
import torch.nn as nn
from copy import deepcopy
from peft import LoraConfig, get_peft_model

class Net(nn.Module):
    def __init__(self):
        super().__init__()
        self.mha = nn.MultiheadAttention(1024, 8)

net = Net()

lora_config = LoraConfig(inference_mode=False,
    r=8,
    lora_alpha=32,
    lora_dropout=0.1,
    target_modules=['mha']
)
lora_net = get_peft_model(deepcopy(net), lora_config)

print(hasattr(net.mha, 'batch_first'))
print(hasattr(lora_net.mha, 'batch_first'))

Output:

True
False

As you can see, batch_first is True for the original MultiheadAttention module, but after converting the model to LoRA, the batch_first attribute is no longer present.

Is this an expected behavior, or is there a way to preserve the batch_first attribute during the conversion?