BenjaminBossan
commented
1 month ago Thanks for proposing this PR to add the Safe LoRA method to PEFT. I have only skimmed the code and paper, so this is not an in-depth review yet.
Based on what I saw, this is my high level understanding. As a user, I start out by training a LoRA adapter on my dataset using an aligned model as the base model. Next, I want to restore safety which may be reduced after training. For this, I take my trained LoRA model, the base model, and the aligned model, Then Safe LoRA will create a new adapter that "injects" the alignment back into my LoRA adapter. Please correct me if my understanding is incorrect.
Implementation-wise, I wonder if we really need a SafeLoRA class. Here is a proposal for a different user API, LMK what you think:
peft_model = ... # my trained LoRA model
# apply_safe_lora is a new function that does what SafeLoRA currently does
# additional arguments like select_layers_type are option args for apply_safe_lora
aligned_model = apply_safe_lora(peft_model, base_model_id, aligned_model_id)
# persist the new safe LoRA adapter
aligned_model.save_pretrained(...)IMO, this would be a simpler API and it would achieve the same thing at the end.
In this example, the aligned_model could actually be the same peft_model as initially, just with a new LoRA adapter loaded which is the Safe LoRA adapter.
Another concern I saw when I read the code is that it looks like it will require a lot of memory. IIUC, we need to have the PEFT model, a copy of the PEFT model, the base model, and the aligned model in memory all at once. Even on CPU RAM, this could be difficult to achieve for many users. I wonder if it would be possible to load the weights from the base model and the aligned model one at a time, at least as an option.
I have already implemented something like this for another use case. Here is the code:
It only works for safetensors, because pickle does not allow to load the weights lazily, but I still think it could be a nice addition.
Also, do we really need a full copy of the PEFT model? Would it be sufficient to only copy the LoRA weights? Ideally, I would like to see a solution that works even if the user has less than twice the memory required for the base model. If this doesn't work, it's also okay, but it would greatly reduce the number of users who can use the method.
chiayi-hsu
github-actions[bot]
Hello,
We have published a paper called Safe LoRA (https://arxiv.org/abs/2405.16833). This work focuses on improving the safety of well-trained LoRA models. Additionally, I have provided an example implementation in the model.py file to illustrate how to apply the Safe LoRA approach effectively.