My goal is to load a pre-trained Hugging Face model, train it, save it, and then load it. Here are the steps that I took:

Loading the Model

from transformers import AutoModelForSequenceClassification, AutoTokenizer, Trainer, TrainingArguments, DataCollatorWithPadding
import numpy as np

# load model
model = AutoModelForSequenceClassification.from_pretrained(
    "gpt2",
    num_labels=6,
    id2label=id2label,
    label2id=label2id,
)

# unfreeze params
for param in model.parameters():
    param.requires_grad = True

# get tokenizer
tokenizer = AutoTokenizer.from_pretrained("gpt2")

# tokenize data
tokenized_ds = {}
for split in ["train", "test"]:
    tokenized_ds[split] = ds[split].map(
        lambda x: tokenizer(x["text"], truncation=True),
        batched=True,
    )

# define an eos token for the tokenizer and model (for padding purposes)
tokenizer.pad_token = tokenizer.eos_token
model.config.pad_token_id = tokenizer.pad_token_id


# define an accuracy metric
def compute_metrics(eval_pred):
    predictions, labels = eval_pred
    predictions = np.argmax(predictions, axis=1)
    return {"accuracy": (predictions == labels).mean()}

PEFTing

from peft import TaskType, LoraConfig, get_peft_model

# create a peft config
config = LoraConfig(
    r=8, 
    lora_alpha=32,
    target_modules=['c_attn', 'c_proj'],
    lora_dropout=0.1,
    bias="none",
    fan_in_fan_out=True,
    task_type=TaskType.SEQ_CLS
)

# create lora model
lora_model = get_peft_model(model, config)
lora_model.print_trainable_parameters()

# unfreeze params
for param in lora_model.parameters():
    param.requires_grad = True

# train model
trainer = Trainer(
    model=lora_model,
    args=TrainingArguments(
        output_dir="./data/lora_peft",
        learning_rate=2e-5,
        per_device_train_batch_size=16,
        per_device_eval_batch_size=16,
        evaluation_strategy="epoch",
        save_strategy="epoch",
        
        num_train_epochs=1,
        weight_decay=0.01,
        load_best_model_at_end=True,
    ),
    train_dataset=tokenized_ds["train"].rename_column('label', 'labels'),
    eval_dataset=tokenized_ds["test"].rename_column('label', 'labels'),
    tokenizer=tokenizer,
    data_collator=DataCollatorWithPadding(tokenizer=tokenizer),
    compute_metrics=compute_metrics,
)

trainer.train()

Saving and Loading (Problem Area)

# save weights to 'gpt-lora'
lora_model.save_pretrained("gpt-lora")

from peft import AutoPeftModelForSequenceClassification

# load and evaluate trained model from 'gpt-lora'
loaded_model = AutoPeftModelForSequenceClassification.from_pretrained("gpt-lora") # <-problem here

The Error

---------------------------------------------------------------------------
RuntimeError                              Traceback (most recent call last)
Cell In[11], line 4
      1 from transformers import AutoModelForSequenceClassification
      3 # load and evaluate trained model
----> 4 loaded_model = AutoModelForSequenceClassification.from_pretrained("gpt-lora")
      5 trainer.evaluate()

File /opt/conda/lib/python3.10/site-packages/transformers/models/auto/auto_factory.py:566, in _BaseAutoModelClass.from_pretrained(cls, pretrained_model_name_or_path, *model_args, **kwargs)
    564 elif type(config) in cls._model_mapping.keys():
    565     model_class = _get_model_class(config, cls._model_mapping)
--> 566     return model_class.from_pretrained(
    567         pretrained_model_name_or_path, *model_args, config=config, **hub_kwargs, **kwargs
    568     )
    569 raise ValueError(
    570     f"Unrecognized configuration class {config.__class__} for this kind of AutoModel: {cls.__name__}.n"
    571     f"Model type should be one of {', '.join(c.__name__ for c in cls._model_mapping.keys())}."
    572 )

File /opt/conda/lib/python3.10/site-packages/transformers/modeling_utils.py:3775, in PreTrainedModel.from_pretrained(cls, pretrained_model_name_or_path, config, cache_dir, ignore_mismatched_sizes, force_download, local_files_only, token, revision, use_safetensors, *model_args, **kwargs)
   3772     model = quantizer.post_init_model(model)
   3774 if _adapter_model_path is not None:
-> 3775     model.load_adapter(
   3776         _adapter_model_path,
   3777         adapter_name=adapter_name,
   3778         token=token,
   3779         adapter_kwargs=adapter_kwargs,
   3780     )
   3782 if output_loading_info:
   3783     if loading_info is None:

File /opt/conda/lib/python3.10/site-packages/transformers/integrations/peft.py:206, in PeftAdapterMixin.load_adapter(self, peft_model_id, adapter_name, revision, token, device_map, max_memory, offload_folder, offload_index, peft_config, adapter_state_dict, adapter_kwargs)
    203     processed_adapter_state_dict[new_key] = value
    205 # Load state dict
--> 206 incompatible_keys = set_peft_model_state_dict(self, processed_adapter_state_dict, adapter_name)
    208 if incompatible_keys is not None:
    209     # check only for unexpected keys
    210     if hasattr(incompatible_keys, "unexpected_keys") and len(incompatible_keys.unexpected_keys) > 0:

File /opt/conda/lib/python3.10/site-packages/peft/utils/save_and_load.py:135, in set_peft_model_state_dict(model, peft_model_state_dict, adapter_name)
    132 else:
    133     raise NotImplementedError
--> 135 load_result = model.load_state_dict(peft_model_state_dict, strict=False)
    136 if config.is_prompt_learning:
    137     model.prompt_encoder[adapter_name].embedding.load_state_dict(
    138         {"weight": peft_model_state_dict["prompt_embeddings"]}, strict=True
    139     )

File /opt/conda/lib/python3.10/site-packages/torch/nn/modules/module.py:2041, in Module.load_state_dict(self, state_dict, strict)
   2036         error_msgs.insert(
   2037             0, 'Missing key(s) in state_dict: {}. '.format(
   2038                 ', '.join('"{}"'.format(k) for k in missing_keys)))
   2040 if len(error_msgs) > 0:
-> 2041     raise RuntimeError('Error(s) in loading state_dict for {}:nt{}'.format(
   2042                        self.__class__.__name__, "nt".join(error_msgs)))
   2043 return _IncompatibleKeys(missing_keys, unexpected_keys)

RuntimeError: Error(s) in loading state_dict for GPT2ForSequenceClassification:
    size mismatch for score.weight: copying a param with shape torch.Size([6, 768]) from checkpoint, the shape in current model is torch.Size([2, 768]).

I keep getting a pytorch size mismatch somewhere in the process. The model seems to save properly.. Just for reference this is the result of saving the model (the first line of code provided above) gpt-lora/adapter_config.json:

{
“auto_mapping”: null,
“base_model_name_or_path”: “gpt2”,
“bias”: “none”,
“fan_in_fan_out”: true,
“inference_mode”: true,
“init_lora_weights”: true,
“layers_pattern”: null,
“layers_to_transform”: null,
“lora_alpha”: 32,
“lora_dropout”: 0.1,
“modules_to_save”: null,
“peft_type”: “LORA”,
“r”: 8,
“revision”: null,
“target_modules”: [
“c_attn”,
“c_proj”
],
“task_type”: “SEQ_CLS”
}