Adds Vera (Vector Based Random Matrix Adaption) #2

[BenjaminBossan]

Continuation of #1039.

Should now be 95% on par with that PR, with some minor changes on my part + resolving merge conflicts.

Examples and docs have not been included yet.

TODOS:

• Add documentation
• Add examples

https://arxiv.org/abs/2310.11454

Notable changes vis-à-vis 1039:

• Some refactors around how the initialization of VeraModel proceeds, should be more straightforward now.
• Add tests around saving and loading, which needs some special considerations for VeRA.
• Fixed some issues with multiple adapters, requires more strictness (e.g. not allowing multiple different prng keys on the same model).
• projection_prng_key now has a valid default value (0) in the config.
• Removed support for Embedding to reduce complexity: Supporting Embedding layers with VeRA makes very little sense because its shape is always different from the linear layers' shapes. Therefore, they cannot share the vera_A and vera_B matrices, resulting in an error. The only conceivable way to support Embedding layers would be to only target that layer (and possibly the output layer if it shares the weight), but that more or less defeats the purpose of using VeRA. We may revisit support for Embeddings in the future, maybe if we can enable vera_A and vera_B to be of different shapes. Until then, let's support the most common use cases and simplify our lives.

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

To ensure that the vera_A and vera_B weights are shared (but not other tensors), I added some tests that check their corresponding data_ptr()s.

Moreover, I wrote a small script to check the amount of memory taken by the model. For this, I used a very high rank of 10000, so that vera_A and vera_B should be quite large. Then I compared the GPU memory taken for a model with a single layer having a VeRA adapter vs a model with many layers having a VeRA adapter. We should expect that both should take roughly the same memory, since most parameters are shared.

Here is the script:

from transformers import AutoModelForCausalLM
from peft import get_peft_model, VeraConfig, LoraConfig
from peft.tuners.vera import VeraLayer
from peft.tuners.lora import LoraLayer
import gc
import torch

RANK = 10000
model_id = "facebook/opt-125m"

config_cls = VeraConfig
layer_cls = VeraLayer

def get_gpu_memory():
    torch.cuda.synchronize()  # Wait for all kernels to finish
    gpu_info = {
        'allocated': f"{torch.cuda.memory_allocated(0) / 2**30:.4f}GB",
        'reserved': f"{torch.cuda.memory_reserved(0) / 2**30:.4f}GB",
    }
    print(gpu_info)

print("before loading the base model")
get_gpu_memory()

model = AutoModelForCausalLM.from_pretrained(model_id).cuda()
print("after loading the model")
get_gpu_memory()

config = config_cls(task_type="CAUSAL_LM", target_modules=["model.decoder.layers.0.self_attn.k_proj"], r=RANK)
model = get_peft_model(model, config)
num_vera_layers = len([m for m in model.modules() if isinstance(m, layer_cls)])
print(f"after adding {num_vera_layers} adapted layers with rank {RANK}")
get_gpu_memory()

del model
torch.cuda.empty_cache()
gc.collect()

print("after resetting")
get_gpu_memory()

model = AutoModelForCausalLM.from_pretrained(model_id).cuda()
print("after loading the base model")
get_gpu_memory()

config = config_cls(task_type="CAUSAL_LM", target_modules=["v_proj", "q_proj"], r=10000)
model = get_peft_model(model, config)
num_vera_layers = len([m for m in model.modules() if isinstance(m, layer_cls)])
print(f"after adding {num_vera_layers} adapted layers with rank {RANK}")
get_gpu_memory()

For VeRA, the results are:

before loading the base model
{'allocated': '0.0000GB', 'reserved': '0.0000GB'}
after loading the model
{'allocated': '0.4677GB', 'reserved': '0.5176GB'}
after adding 1 adapted layers with rank 10000
{'allocated': '0.5264GB', 'reserved': '0.5762GB'}
after resetting
{'allocated': '0.0000GB', 'reserved': '0.0000GB'}
after loading the base model
{'allocated': '0.4677GB', 'reserved': '0.5176GB'}
after adding 24 adapted layers with rank 10000
{'allocated': '0.5273GB', 'reserved': '0.5762GB'}

As we can see, when adapting 24 layers vs 1 layer, the memory used is almost identical. We expect a small increase because vera_lambda_b and vera_lambda_d are not shared, so this is in line with our expectations.

As a sanity check, if we do the same with LoRA instead of VeRA, we see a big increase in memory used:

before loading the base model
{'allocated': '0.0000GB', 'reserved': '0.0000GB'}
after loading the model
{'allocated': '0.4677GB', 'reserved': '0.5176GB'}
after adding 1 adapted layers with rank 10000
{'allocated': '0.5263GB', 'reserved': '0.5762GB'}
after resetting
{'allocated': '0.0000GB', 'reserved': '0.0000GB'}
after loading the base model
{'allocated': '0.4677GB', 'reserved': '0.5176GB'}
after adding 24 adapted layers with rank 10000
{'allocated': '1.8740GB', 'reserved': '1.9238GB'}

All if this is a strong indicator to me that the memory sharing actually works. If anyone has ideas for more tests, let me know.

[BenjaminBossan]

@dkopi @vvvm23 I think I'm pretty much finished with the implementation itself, docs and examples are yet to come. Still, if you have time, I'd be happy with a review or if you can run some tests to see if the implementation performs as expected. The changes compared to the original PR are documented above, the core VeRA computation hasn't been changed, though.

[review-notebook-app]

Check out this pull request on 

See visual diffs & provide feedback on Jupyter Notebooks.

---

Powered by ReviewNB

[vvvm23]

Hi @BenjaminBossan, I can do a review some time this week.

[dkopi]

Looks good 👌

[BenjaminBossan]

@vvvm23 Did you have time to take a look?

[vvvm23]

looks good to me, few small nitpicks. sorry for the delay on this!

[vvvm23]

[nit] why not raise NotImplementedError? Avoid silent failures if something incorrectly calls the hook.

[BenjaminBossan]

Passing is a valid outcome here, if we raised here, all non-VeRA adapters would suddenly error ;)

[vvvm23]

perhaps we should increase the default value? 8 is rather small for VeRA (paper used 256-1024 for their experiments)

[BenjaminBossan]

Yes, makes sense, I'll go with 256.

[vvvm23]

0.1 may be a better default value, see Table 6 in the paper

[BenjaminBossan]

Right, makes sense.

[vvvm23]

remove this TODO? I feel this behaviour is actually fine, the semantics of Conv1D are unlikely to change

[BenjaminBossan]

Done.

[BenjaminBossan]

Thanks for the feedback, Alex, your comments should be addressed now.

[BenjaminBossan]

Passing is a valid outcome here, if we raised here, all non-VeRA adapters would suddenly error ;)

[BenjaminBossan]

Yes, makes sense, I'll go with 256.

[BenjaminBossan]

Right, makes sense.

[BenjaminBossan]

Done.

[pacman100]

Thank you @BenjaminBossan for all the work on Vera continuing the efforts of @vvvm23, all looks great with examples, documentation and tests! 🔥🚀✨

It would be great to add @vvvm23 and @dkopi as co-authors for all their guidance and work!

Left a minor nit.

[pacman100]

Don't think all the target modules have same shape and this also includes embedding layer.

[BenjaminBossan]

@pacman100 Thanks for the feedback, indeed I hadn't checked the docstring example. It is now changed to a working model.

Your comment also prompted me to take a look at the models that are pre-configured in TRANSFORMERS_MODELS_TO_VERA_TARGET_MODULES_MAPPING. This was just a copy of the the LoRA settings. Unfortunately, not all models work, I had to exclude some popular ones like Mistral, Mixtral, Phi, and gemma. The issue is again that the shapes of the target layer weights can differ. Hopefully, we can add the feature of supporting multiple different weight shapes in the future.

It would be great to add @vvvm23 and @dkopi as co-authors for all their guidance and work!

Yes, that was indeed my plan. @vvvm23 @dkopi could you please let me know how you want to be added as co-authors?

[dkopi]

@BenjaminBossan You can add:
Co-authored-by: Dawid <[email protected]>
Thanks!

[vvvm23]

Likewise, you can add
Co-authored-by: Alex McKinney <[email protected]>

Thanks @BenjaminBossan for bringing this PR to completion!

[BenjaminBossan]

Done 🎉

Thanks again so much @vvvm23 for doing the majority of the work and @dkopi for your constant feedback.

Let's hope that VeRA gains traction in the community. For the future, I'll add this list of improvements for VeRA that have yet to be implemented (contributions are welcome):

1. Make VeRA work with different weight shapes. This is IMO the biggest limitation right now. Most straightforward way would be to have one pair of fixed vera_A/vera_B weights per target weight shape.
   There are cases where the shapes are the same when transposed, so (4096, 1024) vs (1024, 4096) (down and up projections for instance) - I wonder if we can use the same weights here and just transpose them.
2. If this is done, implement VeRA for more layer types than just Linear. Right now, supporting, for instance, Embedding, makes little sense, as it almost always has a different shape than Linear.
3. Support quantized weights, most notably bnb.
4. Support DoRA ("DVoRA")

[vvvm23]

Thanks again @BenjaminBossan ! please tag me in issues and PRs related to improvements :)