Missing trainable parameters in a loaded LoRA model

🤗Transformers
1

I followed this blog post for saving and loading a PEFT model. I defined my model as

model = AutoModelForSequenceClassification.from_pretrained('roberta-base', return_dict=True)
peft_config = LoraConfig(task_type="SEQ_CLS", inference_mode=False, r=8, lora_alpha=16, lora_dropout=0.1)
model = get_peft_model(model, peft_config)
model.print_trainable_parameters()

I saved it with

save_path = './test_save'
model.save_pretrained(save_path)

and load it with

config = PeftConfig.from_pretrained(save_path)
model = AutoModelForSequenceClassification.from_pretrained(config.base_model_name_or_path)
model = PeftModel.from_pretrained(model, save_path)
model.print_trainable_parameters()

However, I got 300k less trainable parameters when loading the same model. When saving the model, it prints

trainable params: 1479172 || all params: 125534212 || trainable%: 1.1783018959006968

after loading, it prints

trainable params: 1184260 || all params: 125534212 || trainable%: 0.9433762964951737

Is this the correct way of saving and loading a PEFT model?

2

Found what’s wrong, model.save_pretrained somehow automatically freezes all LoRA parameters and only keep the last one or two fully connected classifier head layers trainable (didn’t see this behaviour documented), so we just need to set all LoRA parameters trainable again with:

for name, param in model.named_parameters():
    if 'lora' in name:
        param.requires_grad = True
2 Likes