Support for Activated-Lora

[vpgits]

Feature request

Source: https://research.ibm.com/blog/inference-friendly-aloras-lora
Github: https://github.com/IBM/activated-lora
Paper: https://arxiv.org/pdf/2504.12397

I would like to propose a solid extension of Peft for this new activated lora implementation. This is relatively new, and there might be no demand for it yet. But I am creating this issue just in case someone is interested.

Motivation

I am planning on using ALora in my undergrad fyp, current implementation from IBM is monkeypatched, and I'd like to have a solid implementation. Therefore I would try to write my own extension for this activated lora implementation.

Your contribution

My technical expertise is quite limited but if I could make good enough progress I would create a PR.

[BenjaminBossan]

Thank you for your proposal to integrate activated LoRA (aLoRA) into PEFT. Also, thanks for offering to work on the PR. I take it you're not one of the paper authors? If you do decide to start, don't hesitate to create an early draft PR so that we can give quick feedback.

I read part of the paper and only very quickly skimmed the code. As is, I'm not sure I fully understand what is involved in using aLoRA. Very naively, I would think we can just keep the existing kv cache when turning on the new LoRA, subsequently filling up the new cache entries with the adapted weight. Since, right now, there is no mechanism in PEFT that would purge the cache when users enable or switch LoRAs, this is basically the default behavior at the moment. However, I have not understood yet what the proposed method does on top of that. If you can summarize the difference, I'd be glad to hear it.

[vpgits]

Thanks for responding at such a short time.

I am not one of the paper's authors. I just found the paper to be interesting. I'll try my best to explain from what I've understood so far but this is kinda out of my league honestly.

Basically they cache the attention layers only (query, key and value projections only), utilize something called late prompting where they first inference with the base model and then reuse when switching to the aLora model.

They are also looking into invoking the aLora model through the base model itself (by generating a invocation string at low confidence responses which will have to be fine tuned however), or providing the name of the intrinsic (they have proposed a role based naming scheme) by appending that to the completion response and support for registering multiple adapters to the same model.

The invocation of the intrinsic (aLora model) is another challenge and correctly identifying the correct position to switch models is challenging, which they have attempted to do through hooks in pytorch by tracking the position of an invocation string token sequence.

There will be some modifications to the forward pass to accommodate these changes.

There was also a data collator in the github repo I've shared, so there's also a need to handle this requirement as well (multi turn conversations with usage of different adapters).

[BenjaminBossan]

Thanks for explaining this further. This indeed sounds like quite a lot of steps. My impression is that this could be very specific to some use cases that matter to IBM (or their customers) and before spending too much time on it and making the code base more complex, we should ideally get an estimate if there is a demand for the feature. Since you found this, maybe you have some insights into applications of the technique.

Also, I'd like to know if this extra effort is measurably better than the very simple approach I mentioned above, but I need to check the paper if they covered this in their experiments.

[vpgits]

import torch,os
from transformers import AutoTokenizer,  AutoModelForCausalLM, DynamicCache
from alora.peft_model_alora import aLoRAPeftModelForCausalLM
from alora.config import aLoraConfig
from alora.tokenize_alora import tokenize_alora

REUSE_CACHE = True  # If True, demonstrate KV cache reuse (slightly more complex code). If False, use simplest generation code.

token = os.getenv("HF_MISTRAL_TOKEN")
BASE_NAME = "ibm-granite/granite-3.2-8b-instruct"
ALORA_NAME = "ibm-granite/granite-3.2-8b-alora-uncertainty"
device=torch.device('cuda' if torch.cuda.is_available() else 'cpu')

# Load model
tokenizer = AutoTokenizer.from_pretrained(BASE_NAME,padding_side='left',trust_remote_code=True, token=token)
model_base = AutoModelForCausalLM.from_pretrained(BASE_NAME,device_map="auto")
model_UQ = aLoRAPeftModelForCausalLM.from_pretrained(model_base, ALORA_NAME)

question = "What is IBM Research?"
print("Question:" + question)
question_chat = [
    {
        "role": "user",
        "content": question
    },
]

# Generate answer with base model
input_text = tokenizer.apply_chat_template(question_chat,tokenize=False,add_generation_prompt=True)
# Remove default system prompt (Granite chat template)
len_sys = len(input_text.split("<|start_of_role|>user"))
input_text = input_text[len_sys:]

#tokenize
inputs = tokenizer(input_text, return_tensors="pt")
if REUSE_CACHE: #save KV cache for future aLoRA call
    prompt_cache = DynamicCache()
    with model_UQ.disable_adapter():
        output_dict = model_base.generate(inputs["input_ids"].to(device), attention_mask=inputs["attention_mask"].to(device), max_new_tokens=600,past_key_values = prompt_cache, return_dict_in_generate=True)
    answer_cache = output_dict.past_key_values
    output = output_dict.sequences
else: #simplest call
    with model_UQ.disable_adapter():
        output = model_UQ.generate(inputs["input_ids"].to(device), attention_mask=inputs["attention_mask"].to(device), max_new_tokens=600)
output_text = tokenizer.decode(output[0])
# Base model answer (split uses Granite chat template)
answer = output_text.split("assistant<|end_of_role|>")[1]
print("Answer: " + answer)

# Generate certainty score
#Get Invocation string to append to input.
uq_generation_prompt = model_UQ.peft_config[model_UQ.active_adapter].invocation_string #For UQ, this is "<|start_of_role|>certainty<|end_of_role|>" 
uq_chat = question_chat + [
    {
        "role": "assistant",
        "content": answer
    },
]

uq_text = tokenizer.apply_chat_template(uq_chat,tokenize=False) 
uq_text = uq_text[len_sys:]
# tokenize and generate
inputs, alora_offsets = tokenize_alora(tokenizer,uq_text, uq_generation_prompt)

if REUSE_CACHE: #reuse KV cache from earlier answer generation 
    output = model_UQ.generate(inputs["input_ids"].to(device), attention_mask=inputs["attention_mask"].to(device), max_new_tokens=1,alora_offsets=alora_offsets,past_key_values=answer_cache)
else: #simplest call
    output = model_UQ.generate(inputs["input_ids"].to(device), attention_mask=inputs["attention_mask"].to(device), max_new_tokens=1,alora_offsets=alora_offsets)
output_text = tokenizer.decode(output[0])

# Extract score
uq_score = int(output_text[-1])
print("Certainty: " + str(5 + uq_score * 10) + "%")

(from https://github.com/IBM/activated-lora/blob/main/experiments/inference_example.py).

As you said, since transformers purges the cache between generations, they use DynamicCache from transformers to cache the attention layers.
The novelty here is the conditional activation of lora, IBM themselves are using this for RAG.

Ps: currently Im trying to do create a draft in my end https://github.com/vpgits/peft/tree/alora .

[vpgits]

Tested out this implementation on https://github.com/vpgits/peft/tree/alora, on jupyter notebook.

!pip uninstall peft -y
!pip install git+https://github.com/vpgits/peft.git@alora
!pip install bitsandbytes

import torch
from transformers import AutoTokenizer

def tokenize_alora(tokenizer, inputs, activation_strings):
    input_tokenized = tokenizer(inputs, return_tensors="pt")
    activation_tokenized = tokenizer(activation_strings, return_tensors="pt")
    input_combined = {
        "input_ids": torch.cat([input_tokenized["input_ids"], activation_tokenized["input_ids"]], dim=1),
        "attention_mask": torch.cat([input_tokenized["attention_mask"], activation_tokenized["attention_mask"]], dim=1),
    }
    alora_offsets = [activation_tokenized["input_ids"].shape[1]-1]
    return input_combined, alora_offsets

import torch, os
from transformers import AutoTokenizer, AutoModelForCausalLM, DynamicCache
from peft import PeftModelForCausalLM
from peft import ALoraConfig

REUSE_CACHE = False
BASE_NAME = "/content/drive/MyDrive/unsloth/granite-3.2-8b-instruct-bnb-4bit"
LORA_NAME = "ibm-granite/granite-uncertainty-3.2-8b-alora"
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

tokenizer = AutoTokenizer.from_pretrained(BASE_NAME, padding_side='left', trust_remote_code=True)
model_base = AutoModelForCausalLM.from_pretrained(BASE_NAME, device_map="auto")
model_UQ = PeftModelForCausalLM.from_pretrained(model_base, LORA_NAME)

question = "What is in the center of the solar system?"
print("Question:" + question)

question_chat = [
    {
        "role": "user",
        "content": question
    },
]

input_text = tokenizer.apply_chat_template(question_chat, tokenize=False, add_generation_prompt=True)
len_sys = len(input_text.split("<|start_of_role|>user"))
input_text = input_text[len_sys:]
inputs = tokenizer(input_text, return_tensors="pt")

if REUSE_CACHE:
    prompt_cache = DynamicCache()
    with model_UQ.disable_adapter():
        output_dict = model_base.generate(
            inputs["input_ids"].to(device), 
            attention_mask=inputs["attention_mask"].to(device), 
            max_new_tokens=600,
            past_key_values=prompt_cache, 
            return_dict_in_generate=True
        )
    answer_cache = output_dict.past_key_values
    output = output_dict.sequences
else:
    with model_UQ.disable_adapter():
        output = model_UQ.generate(
            inputs["input_ids"].to(device), 
            attention_mask=inputs["attention_mask"].to(device), 
            max_new_tokens=600
        )

output_text = tokenizer.decode(output[0])
answer = output_text.split("assistant<|end_of_role|>")[1]
print("Answer: " + answer)

uq_generation_prompt = "<|start_of_role|>certainty<|end_of_role|>"
uq_chat = question_chat + [
    {
        "role": "assistant",
        "content": answer
    },
]

uq_text = tokenizer.apply_chat_template(uq_chat, tokenize=False)
uq_text = uq_text[len_sys:]
inputs, alora_offsets = tokenize_alora(tokenizer, uq_text, uq_generation_prompt)

if REUSE_CACHE:
    output = model_UQ.generate(
        inputs["input_ids"].to(device), 
        attention_mask=inputs["attention_mask"].to(device), 
        max_new_tokens=1,
        alora_offsets=alora_offsets,
        past_key_values=answer_cache
    )
else:
    output = model_UQ.generate(
        inputs["input_ids"].to(device), 
        attention_mask=inputs["attention_mask"].to(device), 
        max_new_tokens=1,
        alora_offsets=alora_offsets
    )

output_text = tokenizer.decode(output[0])
uq_score = int(output_text[-1])
print("Certainty: " + str(5 + uq_score * 10) + "%")

Output from the generation segment was

Answer: The center of the solar system is occupied by the Sun. It's a massive, luminous ball of hot plasma that contains about 99.86% of the mass of the entire solar system. The planets, including Earth, orbit around the Sun due to its immense gravitational pull.<|end_of_text|>
Certainty: 85%

I think this state of the branch can be considered as a draft implementation.

[BenjaminBossan]

Thank you for working on this. I think for us, it would be easiest if you opened a draft PR so that we can see what additions to PEFT are necessary.

[kgreenewald]

Hi everyone - thanks for starting this thread, and thanks to vpgits for your initiative here! I am the main author of the paper and the github repo. We are actively working towards putting together a PR for alora integration into peft.

In the meantime, some answers to questions raised:

"Very naively, I would think we can just keep the existing kv cache when turning on the new LoRA, subsequently filling up the new cache entries with the adapted weight. Since, right now, there is no mechanism in PEFT that would purge the cache when users enable or switch LoRAs, this is basically the default behavior at the moment."

-> There's a subtle yet crucial misunderstanding here: everything you say here is true but that's not what we are doing. Activated lora is not about "keeping around existing kv cache", but instead about reusing KV cache that was created by a prior call to the base model.

Think of it this way.

Not yet worrying about adapters: If you have a single model, KV cache reuse is possible when you have multiturn interactions. User turn, generate an assistant turn, new user turn, new generated assistant turn. Here, on the second model call, the KV cache created during the first model call can be reused, saving some prefill time.

Now, consider a multi-model flow. E.g. User turn, base model generates assistant turn, new user turn, adapter model generates assistant turn. Naively, if the adapter here is LoRA, the second call cannot reuse the KV cache that the base model created in the first call, because the LoRA has different model weights. Hence, the prior turns need to be prefilled into the LoRA.

Activated LoRA solves this issue by only adapting weights for tokens coming after the token at which the adapter is invoked.

Importantly, this is not about a different way to do inference with existing LoRAs, this requires a separate training run with the modified architecture.

That said, I have made the aLoRA code to conform as close as possible to the LoRA code in PEFT, so the PR should not be too invasive and the training and inference code will appear very similar between LoRA and aLoRA.

Another point here: most of these efficiency gains will be seen at runtime on serving platforms like vLLM. We're working on that separately. Where PEFT comes in is defining and training these adapters using this new architecture (e.g. SFTTrainer).

---

"This indeed sounds like quite a lot of steps. My impression is that this could be very specific to some use cases that matter to IBM (or their customers) and before spending too much time on it and making the code base more complex, we should ideally get an estimate if there is a demand for the feature. Since you found this, maybe you have some insights into applications of the technique."

I hope the above clarifies this as well - the idea is actually quite simple despite the complexity of the code demonstrating the cache reuse that was cited. This complexity is mostly a function of how HF does the DynamicCache, separately, we are also working on a PR to vLLM that is much more seamless, with all cache reuse happening automatically without needing to pass cache objects around.

Happy to discuss more if the above isn't clear! We'll be moving towards a PR soon.

[githubnemo]

Thanks for the insight @kgreenewald and glad to hear you're working on a PR.

In that light, I think it would make sense to not proceed with #2530 for the time being to avoid doing the same work twice. WDYT @vpgits?

[vpgits]

Yes, I have stopped working on #2530 on my end. Thanks for reaching out.

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

[githubnemo]

not stale