AdaLora rank_pattern is None, Docs and Warnings and Integrations should be improved up on.

[kuronekosaiko]

System Info

docker.io/pytorch/pytorch:2.1.2-cuda12.1-cudnn8-runtime
Python 3.10.13
pytorch 2.1.2
peft 0.7.1
transformers 4.36.2
bitsandbytes 0.41.3.post2

Who can help?

@pacman100 @younesbelkada @sayakpaul

Information

• The official example scripts
• My own modified scripts

Tasks

• An officially supported task in the examples folder
• My own task or dataset (give details below)

Reproduction

from mylib import my_custom_dataset
from transformers import GPTQConfig
import torch
from transformers import (
    AutoModelForCausalLM,
    AutoTokenizer,
    Trainer,
    TrainingArguments
)
from peft import (
    prepare_model_for_kbit_training,
    PeftModelForCausalLM,
    AdaLoraConfig
)


# Yi is a llama family model.
model_path = "models/Yi-6B-GPTQ"
dataset = my_custom_dataset()

quantization_config = GPTQConfig(bits=4, use_exllama=False)
model = AutoModelForCausalLM.from_pretrained(model_path, quantization_config=quantization_config)

model.gradient_checkpointing_enable()
model = prepare_model_for_kbit_training(model)

config = AdaLoraConfig(
    lora_alpha=8,
    target_r=4,
    init_r=24,
    total_step=20,
    tinit=5,
    tfinal=2,
    deltaT=1,
    lora_dropout=0,
    task_type="CAUSAL_LM",
    bias="none",
    target_modules=[
        "q_proj",
        "k_proj",
        "v_proj",
        "o_proj",
        "down_proj",
        "up_proj",
        "gate_proj"
    ]
)
model = PeftModelForCausalLM(model, config, "default")

model.config.use_cache = False
trainer = Trainer(
    model=model,
    train_dataset=dataset,
    args=TrainingArguments(
        per_device_train_batch_size=1,
        gradient_accumulation_steps=8,
        max_steps=20,
        warmup_steps=5,
        logging_steps=1,
        save_steps=2,
        learning_rate=2**(-12),
        gradient_checkpointing=True,
        fp16=True,
        fp16_opt_level="O2",
        output_dir="adalora_test5_outputs",
        logging_strategy="steps",
        lr_scheduler_type="constant_with_warmup",
        save_strategy="steps",
        save_total_limit=2,
        optim="paged_adamw_32bit"
    )
)

with torch.backends.cuda.sdp_kernel(enable_flash=False,
                                    enable_math=False,
                                    enable_mem_efficient=True):
    trainer.train()

Expected behavior

rank_pattern in adapter_config.json should be not None/null. Something like this, a trained adalora adapter, a not empty rank_pattern is expected.

rank_pattern should be updated in every DeltaT step, and the code should update the "Adaptive Budget Allocation" mask using this rank_pattern in every DeltaT step.

Since ~8 months ago as shown in this, rank_pattern is no longer correct.

From my testing (code shown below), neither of AdaLoraModel.update_and_allocate and RankAllocator.update_and_allocate are ever called. This might be the reason of the above behaviour.

......

from peft import AdaLoraModel
from peft.tuners.adalora.layer import RankAllocator

# a function to raise exception if ever called
def raise_anyway(self, *args, **kwargs):
    raise

# monkey-patch
AdaLoraModel.update_and_allocate = raise_anyway
RankAllocator.update_and_allocate = raise_anyway

# load the model and train as usual without any errors.
......

[kuronekosaiko]

After further inspection, the update_and_allocatemethod should be called manually in every training step just before zero_grad, so it won't work with huggingface's trainer.

As far as I know, adalora is the only officially supported peft method that have this problem.

However this quirk is not communicated to users, either through docs or warnings, and because the integration is so similar to the other peft methods, this is prone to be overlooked.

And in fact most adalora adapter uploaded to the huggingface hub have "rank_pattern": null in their adapter_config.json, tells this is exactly the case.

My suggestion is:

1. Update Docs, add update_and_allocate method description and examples to AdaLoraModel class in the doc. So that this quirk is communicated clearly.
2. Warn users, display a warning that says "Adalora in its current state does not work with trainer, because of update_and_allocate.", either when using PeftModelForCausalLM/get_peft_model, or when loading huggingface's trainer.
3. improve integration, with something like a prepare_trainer_for_adalora_training function similar to how prepare_model_for_kbit_training works, to add update_and_allocate to the training loop.

[BenjaminBossan]

Thanks for reporting this issue. Indeed, update_and_allocate needs to be called explicitly during the training loop, as shown for instance here. I agree with your first point, that this should be highlighted in the docs. Are you interested in working on a PR to do so?

Regarding your second point, I would not go so far as to add a warning. Many users don't use Trainer but some other framework or write their training loop themselves. For those, the warning would be superfluous. We should reduce the number of warnings that are not strictly necessary, otherwise users start to ignore them.

Regarding your third point, I'm not really sure how that should work, as PEFT has no control over the training loop. In theory, such a functionality could be added to Trainer, but this is far too specific and I don't think it has a chance to be accepted. However, it should be possible to subclass Trainer or write a custom callback to achieve this functionality, maybe show them in the docs too. I'm not very experienced with Trainer, though, so someone else would have to do it.

Regarding rank_pattern, this is a bit of a different story. It is not intended to hold the ranks derived by AdaLora. Instead, its intent is to initialize different LoRA layers with different ranks. Therefore, this argument should not be updated by AdaLora.

[kuronekosaiko]

I agree what you said, but

Regarding rank_pattern, this is a bit of a different story. It is not intended to hold the ranks derived by AdaLora. Instead, its intent is to initialize different LoRA layers with different ranks. Therefore, this argument should not be updated by AdaLora.

About this, as shown in the below code block, update_and_allocate would store its mask information in peft_config.rank_pattern. Due to lora_config = self.peft_config[self.trainable_adapter_name] and lora_config.rank_pattern = rank_pattern, it's python being python.

peft/src/peft/tuners/adalora/model.py

Lines 307 to 327 in e96eef9

	def update_and_allocate(self, global_step):
	lora_config = self.peft_config[self.trainable_adapter_name]
	# Update the importance score and allocate the budget
	if global_step < lora_config.total_step - lora_config.tfinal:
	_, rank_pattern = self.rankallocator.update_and_allocate(self.model, global_step)
	if rank_pattern:
	lora_config.rank_pattern = rank_pattern
	# Finalize the budget allocation
	elif global_step == lora_config.total_step - lora_config.tfinal:
	_, rank_pattern = self.rankallocator.update_and_allocate(self.model, global_step, force_mask=True)
	# for some reason, this freezes the trainable parameters and nothing gets updates
	# self.resize_modules_by_rank_pattern(rank_pattern, self.trainable_adapter_name)
	lora_config.rank_pattern = rank_pattern
	self.rankallocator.reset_ipt()
	# Currently using inefficient way to mask the unimportant weights using the rank pattern
	# due to problem mentioned above
	elif global_step > lora_config.total_step - lora_config.tfinal:
	self.rankallocator.mask_using_rank_pattern(self.model, lora_config.rank_pattern)
	# Pass the function and do forward propagation
	else:
	return None

And I was wondering that too, about why storing mask information in rank_pattern, which gives rank_pattern double duty. If it is not intended, then it might need fixes. However multiple lines of code is contributing and relying on this behavior, and I think the mask information are needed when loading a trained adalora adapter.

[kuronekosaiko]

However, it should be possible to subclass Trainer or write a custom callback to achieve this functionality, maybe show them in the docs too.

About this, I believe it's possible to do with "inject custom behavior" mentioned in huggingface trainer's doc, but I'm not sure.

The Trainer contains the basic training loop which supports the above features. To inject custom behavior you can subclass them and override the following methods:
......
training_step
......

Something like this would be preferable:

def prepare_trainer_for_adalora_training(trainer):
    import types

    def adalora_training_step():
        # a modified training_step for trainer.

    trainer.training_step = types.MethodType(adalora_training_step, trainer)

trainer = Trainer(model=model, train_dataset=train_dataset, args=args)
prepare_trainer_for_adalora_training(trainer)

[BenjaminBossan]

And I was wondering that too, about why storing mask information in rank_pattern, which gives rank_pattern double duty. If it is not intended, then it might need fixes. However multiple lines of code is contributing and relying on this behavior, and I think the mask information are needed when loading a trained adalora adapter.

Oh wow, I totally missed that, thanks for pointing me to it. I guess it is intentional that rank_pattern performs double duty here, but @pacman100 would have to confirm this, as he implemented rank_pattern for LoRA in #873.

Something like this would be preferable:

Not sure why this is preferable to subclassing and overriding training_step. Subclassing seems to be cleaner to me.

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