(PEFT hugginFace) Fine tuning a pre-trained model with additional layers on top (subclassed pre-trained model)

[Lpourchot]

Hello,

I'm wonder how (and if it is possible) to use Hugginface PEFT fine tuning with a subclassed pre-trained model with new layers.
Example :

    class My_New_Model(torch.nn.Module):
        def __init__(self):
              super(My_New_Model, self).__init__()
              self.bert_model = bert_model              # Pre-trained model
              self.linear = torch.nn.Linear(768, 2)      # Target : 2 values, for regression multi labels for example
    
    def forward(self, input_ids, attention_mask, token_type_ids):
             output = self.bert_model(input_ids, 
                                 attention_mask = attention_mask, 
                                 token_type_ids = token_type_ids)
             output = self.linear(output.pooler_output)
             return output

[BenjaminBossan]

It depends on the fine-tuning technique, but yes, in general it should work. E.g. for LoRA, you can set the target_modules you want to fine-tune in the config. For the new linear layer, you might want to add it to modules_to_save.

[Lpourchot]

Hi Benjamin, Thank you.
Just in case, other people would be interested :
I applied LoRA on my model : Bert pre-trained model with custom layers on top. LoRA added layers for training and froze others (as expected). So I found easier to directly "re-enable" my custom layers for training, and trained the model as usual pytorch model (model.train...) :

lora_model = get_peft_model(my_model, config_LoRA)
lora_model.model.my_layer.weight.requires_grad = True

I trained LoRA and my custom layers and it works fine :-)

[BenjaminBossan]

Great that it works. Your solution is valid for training, but just be cautious that if you use lora_model.save_pretrained, my_layer will not be saved, since peft does not "know" about my_layer. This is why adding it to modules_to_save in the config is preferred.

[Lpourchot]

Thank you for the reminder. I saved the model with usual torch.save(model.state_dict()). For predictions within a new session , I reload the model (of course) and the weight dictionnary (model.load_state_dict...)