LORA - how to determine what module_to_save

[SuperBowser]

I am reading through LORA tutorial and one of the options in LoraConfig is modue_to_save. In the example its value is ‘decode-head’. I would like to use LORA with SequenceClassification model and I not sure what module I need to save.

Any thoughts?

[John6666]

If you specify task_type, PEFT will automatically set module_to_save to an appropriate value. If you want to manually search for the head module to save, it would look something like this.

from transformers import AutoModelForSequenceClassification, AutoConfig
import torch.nn as nn

HEAD_CANDIDATES = ("classifier", "score", "logits_proj", "classification_head")

def find_cls_head_name(model):
    present = [n for n, _ in model.named_modules() if n.split(".")[-1] in HEAD_CANDIDATES]
    if present: return present[0], present
    num_labels = getattr(getattr(model, "config", object()), "num_labels", None)
    hits = []
    for parent_name, module in model.named_modules():
        for child_name, child in module.named_children():
            if isinstance(child, nn.Linear) and getattr(child, "out_features", None) == num_labels:
                hits.append(child_name if parent_name == "" else f"{parent_name}.{child_name}")
    return (hits[0] if hits else None), hits

def print_head_name(model_name):
    cfg = AutoConfig.from_pretrained(model_name)
    model = AutoModelForSequenceClassification.from_pretrained(model_name, config=cfg)
    best, all_hits = find_cls_head_name(model)
    print("Model name:", model_name)
    print("All candidate heads:", all_hits)
    print("Suggested modules_to_save:", [best] if best else None)

print_head_name("distilbert-base-uncased-finetuned-sst-2-english")
#Model name: distilbert-base-uncased-finetuned-sst-2-english
#All candidate heads: ['classifier']
#Suggested modules_to_save: ['classifier']
print_head_name("HuggingFaceTB/SmolLM-135M")
#Model name: HuggingFaceTB/SmolLM-135M
#All candidate heads: ['score']
#Suggested modules_to_save: ['score']