[LoRA] Adds support for bias in LoRA #5733
Conversation
|
@Yard1 Thanks for reviewing the PR. I have added the enable_lora_bias flag (default set to false), which prevents the allocation of lora bias tensors when false. |
|
Related: #5930 |
|
@Yard1 Thanks for reviewing. I've added an e2e test for the lora_bias support. |
|
@followumesh you need to run |
|
@followumesh apologies, this needs another conflict resolution! |
| @@ -64,6 +64,64 @@ def dec(*args, **kwargs): | |||
| return dec | |||
|
|
|||
|
|
|||
| def apply_bias( | |||
There was a problem hiding this comment.
Would it be possible to add this inside PunicaWrapper.add_lora()? There could be an optional bias argument in add_lora() and then the logic of testing whether the bias is None and doing the index computation could be moved inside this function. It seems to me that it would eliminate repeated code lines in this file, but I don't know all the details.
There was a problem hiding this comment.
@maxdebayser I don't have a strong opinion, but this was to keep the changes out of punica wrapper because of it not being directly releated to punica.
vllm/lora/utils.py
Outdated
| assert parts[0] == "base_model" | ||
| assert parts[1] == "model" | ||
| if parts[-1] == "weight": | ||
| assert parts[-2] == "lora_A" or parts[-2] == "lora_B" |
There was a problem hiding this comment.
This assertion is failing a couple of lora_modules that we have:
>>> lora_modules[0].split(".")
['base_model', 'model', 'lm_head', 'weight']
>>> lora_modules[1].split(".")
['base_model', 'model', 'model', 'embed_tokens', 'weight']
Still investigating if this would have thrown an error with the previous code as well...
There was a problem hiding this comment.
@followumesh actually it looks like this was a bad merge .. it's reverting a recent change that was made to improve the error message
There was a problem hiding this comment.
@prashantgupta24 Can you point to a lora module I can test with?
There was a problem hiding this comment.
Let me try to find one, also reverting this change gives me the error
ValueError: base_model.model.lm_head.weight is unsupported LoRA weight
There was a problem hiding this comment.
To summarize, since vLLM expects only LoRA weights in the safetensors file, it was actually an error in our lora adapter that we had lm_head within it. But I think the original error is now being reverted because of this change. Ideally, the error should be
ValueError: base_model.model.lm_head.weight is unsupported LoRA weight
Instead, an assertion error assert parts[-2] == "lora_A" or parts[-2] == "lora_B" is now being thrown, which doesn't provide the right detail
There was a problem hiding this comment.
Yes, @followumesh I think the changes here are not actually related to the main PR changes, wonder if you could revert (and any other changes in same category)?
There was a problem hiding this comment.
@prashantgupta24 Can you check now?
There was a problem hiding this comment.
Yeah this looks better. I'll let you know when I get a chance to test it!
Motivation
PEFT, https://github.com/foundation-model-stack/fms-hf-tuning includes support for tuning LoRA bias. This PR enables bias for lora, so the adapters with bias will work with vLLM.
Changes Included