Possible to build a LoRA that doesn't inject into the transformer?

[AngledLuffa]

Feature request

Is it possible to build a LoRA that doesn't inject into the transformer? This would allow for reusing the same basic transformer with multiple adapters in the same process while saving on GPU memory (probably at the expense of some speed)

Motivation

We've started using PEFT with LoRA for tasks such as sentiment analysis and constituency parsing in Stanza, and one thing we found is that there is currently no memory savings compared to using a fully finetuned transformer.

For example, if the transformer loaded for sentiment analysis takes 3GB, with no finetuning we can reuse the same transformer weights when constituency parsing, making for a total of 3GB plus the prediction heads of the models. If we use fully FT transformers, obviously that increases to 6GB assuming those are our only two tasks.

PEFT with LoRA uses inject_adapter_in_model to update the model with the As and Bs, AFAIK, meaning that loading those two models still takes 6GB. If we could have a version of the transformer which does inference with the As and Bs not injected, but wrapping the base transformer's tensors, this would almost certainly be noticeably slower but would allow for a much smaller memory footprint.

Thanks for the extremely useful library, BTW

Your contribution

I probably don't have much time to investigate this in the next couple months, but in the long term it is something I could attempt with some guidance on where to look

[BenjaminBossan]

It would be possible to build a LoRA adapter that works purely through forward hooks on the base model, but that's a big difference to how we implement this right now and also not as flexible.

Regarding your issue, do you know that you can disable the LoRA-adapters completely and then the model behaves like the pure base model? This should allow you to avoid loading the model twice:

with model.disable_adapter():
    # do inference on base model

[AngledLuffa]

Regarding your issue, do you know that you can disable the LoRA-adapters completely and then the model behaves like the pure base model?

Interesting, I did not know that. However, wouldn’t that give different weights to the prediction head, probably resulting in errors since it wasn't trained to recognize those weights?

What I was hoping for was something where each addition use of peft lora only incurred a space cost of |A|+|B|, the same size as the saved weights, but it sounds like such a thing currently doesn't exist, is that correct? If it makes it any easier, I would not be expecting to train from that state, just run inference.

Thanks for the fast reply!

[BenjaminBossan]

wouldn’t that give different weights to the prediction head, probably resulting in errors since it wasn't trained to recognize those weights?

I think I don't quite get your problem yet. My initial understanding was that you need to have the base model for inference and the LoRA-augmented model for inference and wanted to avoid loading the base model twice. Are the prediction heads trained with or without LoRA? If you train them with LoRA, you have probably added them to modules_to_save, which wraps them in a ModulesToSaveWrapper. This one also works with `disable_adapter'.

What I was hoping for was something where each addition use of peft lora only incurred a space cost of |A|+|B|, the same size as the saved weights, but it sounds like such a thing currently doesn't exist, is that correct?

I'm not sure I understand. When you load the base model, it takes 3GB as you mentioned. When you load LoRA on top, that should increase memory only by a tiny amount. My earlier understanding was that the additional memory comes from loading the base model twice.

[AngledLuffa]

I'm not sure I understand. When you load the base model, it takes 3GB as you mentioned. When you load LoRA on top, that should increase memory only by a tiny amount.

We might be misunderstanding on our end, but our belief was the underlying model gets changed by loading LoRA. So if we load two separate peft models on top of the same transformer, the two loadings of peft models will clobber each other:

    from transformers import AutoModel
    from peft import LoraConfig, get_peft_model, set_peft_model_state_dict

    original_bert_model = AutoModel.from_pretrained(model_name)
    sentiment_bert_model = get_peft_model(original_bert_model, sentiment_peft_config)
    set_peft_model_state_dict(sentiment_bert_model, sentiment_lora_state_dict)

    constituency_bert_model = get_peft_model(original_bert_model, constituency_peft_config)
    set_peft_model_state_dict(constituency_bert_model, constituency_lora_state_dict)

We checked that the weights in original_bert_model were changed after calling set_peft_model_state_dict the first time, which led us to believe that the models would overwrite each other in such a code snippet

[BenjaminBossan]

LoRA adapters do indeed mutate the base model, but it doesn't "clobber" it. We took care of making it so that the base model itself can still make predictions as if LoRA was not there when LoRA is being disabled, and different LoRA adapters do play nicely with each other.

Note, however, that I would load the adapter a bit differently than you do, not sure if your code would work correctly:

from peft import PeftModel

base_model = AutoModel...
# load adapter 0, it is automatically the active adapter
peft_model = PeftModel.from_pretrained(base_model, <path-to-adapter0>, adapter_name=<name0>)
# load adapter 1, adapter 0 is still active
peft_model.load_adapter(<path-to-adapter1>, adapter_name=<name1>)
# activate adapter 1, deactivate adapter 0
peft_model.set_adapter(<name1>)

[AngledLuffa]

We took care of making it so that the base model itself can still make predictions as if LoRA was not there when LoRA is being disabled

The bolded part is a little worrisome - does that mean that base_model in my example won't get back the same results unless the peft models are explicitly disabled? We have been treating the annotators such as POS, sentiment, parser, etc as separate entities which had no idea there were other annotators out there that might affect the transformer they are using.

I can verify this myself later tonight, although I don't have time to do so this early afternoon.

[AngledLuffa]

Can confirm there is an issue with our software loading the transformer and then having that transformer overwritten when connecting with LoRA using peft.

We have a POS model and a sentiment model which both use electra-large, except the sentiment model uses a peft wrapper to get better results. We weren't able to figure out how to get better POS tags with peft, but that's a story for another day.

If I load the POS model by itself, it gets an overall score of 93.92 on the EWT UD dataset (those exact details aren't super relevant).

If I load both the POS model and the Sentiment model, using two different invocations of electra-large, it still gets that score.

If instead I load the POS model, then reuse the electra-large transformer to load the Sentiment model, the score drops to 92.12. Our current release uses the first scheme to keep the results consistent, but that results in 6GB of transformer instead of 3GB.

In terms of separation of concerns, I'd love to have a mechanism where the POS model can load its transformer, the Sentiment model can reuse the same transformer and wrap it in peft, and the two models wouldn't have to do anything specific to their transformer before using it.

[BenjaminBossan]

Just so I understand correctly, the same base model (electra large) is used by the POS model and the sentiment model. You load LoRA weights on top of the sentiment model. In that case, indeed, it is expected that the POS model also behaves differently because PEFT will modify the base model by attaching the adapters.

If you want to avoid having a copy of the base model in memory, I would recommend checking if disabling the adapters on the sentiment model restores the performance of the POS model.

[AngledLuffa]

It's actually giving me an error, saying that there is no adapter loaded:

  File "/home/john/stanza/stanza/pipeline/sentiment_processor.py", line 55, in _set_up_model
    trainer = Trainer.load(filename=filename,
  File "/home/john/stanza/stanza/models/classifiers/trainer.py", line 176, in load
    model.bert_model.disable_adapters()
  File "/usr/local/lib/python3.9/site-packages/transformers/integrations/peft.py", line 322, in disable_adapters
    raise ValueError("No adapter loaded. Please load an adapter first.")

I saw your earlier suggestion of using PeftModel.from_pretrained and peft_model.load_adapter as a mechanism for creating peft models from the transformers, but currently that would be a pretty invasive change on our end, since we have been shipping single files with both the model head and the finetuned transformer / peft weights. Hence the use of set_peft_model_state_dict

[BenjaminBossan]

Did you call disable_adapters after you loaded the PEFT adapter? If yes, it would be strange for that error to occur and I would assume that it is somehow related to how you load the adapter.

Also, as I'm reading your code again, how does the constituency_bert_model fit the picture? Or is that unrelated?

[AngledLuffa]

Yes, definitely. This crash happens if I do this:

             set_peft_model_state_dict(model.bert_model, model_params['bert_lora'])
            model.bert_model.disable_adapters()

I'm not giving it a name as part of the set... call. Does that affect whether or not the model thinks it has an adapter attached to it? As I mentioned earlier, I was using get_peft_model to build the peft-wrapped transformer. I tried this instead:

model.bert_model.add_adapter(lora_config, adapter_name="adapter_1")
model.bert_model.disable_adapters()

This works, actually, and we get back the expected scores! This is a bit challenging to work with, though, for a couple reasons.

The simplest problem is, what if we have two Sentiment models or other model that uses a peft adapter? We've tried it with POS, although we haven't yet found settings that consistently improve the POS scores, and there are definitely situations where we need multiple POS models built from the same transformer. They can't all be named adapter_name="pos"...

There's also the problem of thread safety and needing to modify the transformer's state (currently active adapter) in order to use it. That might actually be a problem if the proposal to remove GIL is implemented in python

I know there's a trick to call copy.deepcopy on a transformer where we could then presumably add separate peft adapters for each model we load, but that does cause an increase in GPU usage. One very satisfactory outcome would be a way to copy the state of the transformer such that there are multiple transformer objects which each use the same underlying weights but know about different sets of adapters. Does such a thing exist?

What I'm getting at is that add_adapter or load_adapter appear to be necessary to organize the adapters on the transformer object, but it doesn't fully address the issue we're running into of wanting to keep each model's state separate from the other model states. If nothing else, we'd run into the same problem we have now as soon as we start loading multiple models of the same type which each want their own adapter.

Also, as I'm reading your code again, how does the constituency_bert_model fit the picture? Or is that unrelated?

constituency_bert_model would be the same transformer with its own adapter. This is actually a great example of the problem of needing multiple adapters for the same model. We found that we get roughly the same results with a peft-adapted transformer as we do the non-peft adapter, but now the model size we distribute is 200MB instead of 1.5GB. So that's a huge improvement, thank you.

A possible issue here is that there are multiple applications for having several constituency models loaded at once, for example using those models in an ensemble to get better results. Again we would run into the limitation of how to name the adapters, such as if the first constituency parser loads its peft adapter with the name "constituency", then the second constituency parser would need to load it with the name "constituency-2" etc. Would that be something where the constituency parser itself knows that there have been N previously loaded models? Would the caller need to keep track of that? Either way, that seems a lot more complicated and less clean than simply having a fresh transformer object (hopefully sharing the underlying weights to save GPU space) which doesn't need to know anything about previously loaded adapters.

[BenjaminBossan]

You can load in multiple adapters and give them any name you want (usually using the add_adapter method in transformers). Then you can choose the active adapter with the set_adapter method.

There's also the problem of thread safety and needing to modify the transformer's state (currently active adapter) in order to use it. That might actually be a problem if the proposal to remove GIL is implemented in python

That's still far in the future and will be opt-in.

[AngledLuffa]

You can load in multiple adapters and give them any name you want (usually using the add_adapter method in transformers).

The thread safety concerns may be further off than we'd like, but the unique names solution doesn't really address the issue of situations where we need more than one POS model or more than one constituency model. We could enforce uniqueness of names somehow, but I was hoping to not add more to the global state than necessary. I suppose a random 10 letter name would almost never repeat and wouldn't require any global state at all...

It would also be a little weird for each model to have the workflow of first turn on their own adapter, then run inference, then turn off their adapter (the latter step being necessary so that any model still using the raw version of the transformer doesn't have to figure out if there even are any adapters to turn off).

Currently what does appear to work for loading time, but not GPU memory, is to load the transformer into memory and then clone it N times for each of the adapters we need. Then we get a situation where one annotator doesn't have to care at all how many other annotators of different types or the same type exist, but at the cost of increased GPU memory usage.

If it's not currently a feature, is it at all feasible to make a wrapper which is just a transformer and a single adapter, where that adapter does not affect inference for other users of the underlying transformer? If that's not anywhere on the current project roadmap, is it something where you'd consider merging a PR that implemented a feature like that?

[BenjaminBossan]

We could enforce uniqueness of names somehow, but I was hoping to not add more to the global state than necessary. I suppose a random 10 letter name would almost never repeat and wouldn't require any global state at all...

I don't get this point. You should know which adapters exist beforehand, so you can just choose some static names like "sentiment", "pos0", "pos1", or not?

It would also be a little weird for each model to have the workflow of first turn on their own adapter, then run inference, then turn off their adapter (the latter step being necessary so that any model still using the raw version of the transformer doesn't have to figure out if there even are any adapters to turn off).

Yes, I agree that it's not super convenient and, if you have to switch each time a new sample comes in, this could add some overhead. As you mention, the alternative would be to have a copy of the model for each adapter in memory. I think the feature we planned to add in #903 would have helped in your situation, but we didn't pursue it further.

If it's not currently a feature, is it at all feasible to make a wrapper which is just a transformer and a single adapter, where that adapter does not affect inference for other users of the underlying transformer? If that's not anywhere on the current project roadmap, is it something where you'd consider merging a PR that implemented a feature like that?

Unfortunately, that's not easy to achieve. The underlying PyTorch model does not have sufficient flexibility (say, via hooks) to add all the features that we need without mutating the model itself. It would probably be possible to have what you asked when focusing on only the subset of features that you need, but you'd have to build that yourself.

[AngledLuffa]

I don't get this point. You should know which adapters exist beforehand, so you can just choose some static names like "sentiment", "pos0", "pos1", or not?

Not necessarily true. We provide a way for people to make an annotation pipeline, and we have no control over how many times the user makes a new pipeline with the same base models without keeping some form of global state.

Still, we did implement a cache system for not excessively loading the same word vectors or transformer too many times. Perhaps that would be a reasonable place to keep track of which annotator names have been used in the past as well.

[switching]... this could add some overhead

How much overhead? Is it rewriting tensors or just flipping a pointer?