Fix resuming PeftModel checkpoints in Trainer 
Author: llohann-speranca
Pull request: yes

# What does this PR do?

Fix an error occurred when Trainer tries to resume a PeftModel from a training checkpoint. That was caused since PeftModel.pre_trained saves only adapter-related data while _load_from_checkpoint expects a saved torch model. This PR fix this issue and allows the adapter checkpoint to be loaded.

Resolves: #24252 

Fixes #24252

## Before submitting
- [ ] This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
- [X] Did you read the [contributor guideline](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md#start-contributing-pull-requests),
      Pull Request section?
- [X] Was this discussed via a Github issue or the [forum](https://discuss.huggingface.co/)? Please add a link
      to it if that's the case. (#24252)
- [ ] Did you make sure to update the documentation with your changes? Here are the
      [documentation guidelines](https://github.com/huggingface/transformers/tree/main/docs), and
      [here are tips on formatting docstrings](https://github.com/huggingface/transformers/tree/main/docs#writing-source-documentation).
- [ ] Did you write any new necessary tests?


## Who can review?

@younesbelkada 

Comments:
_The documentation is not available anymore as the PR was closed or merged._

Although this continues training but it doesnt retain the old stuff for me. Can someone look into this?


This doesn't seem like older resume from checkpoint that has it for pytorch models. Inside trainer.train we need to pass resume from checkpoint parameters as our last checkpoint path. while passing the path, it shows as can't find a valid checkpoint. Could someone please post some code snippet on how to use resume from checkpoint for PEFT models?

@pacman100 Requesting to review and merge the changes, thanks!

Hi @techthiyanes 

> This doesn't seem like older resume from checkpoint that has it for pytorch models. Inside trainer.train we need to pass resume from checkpoint parameters as our last checkpoint path. while passing the path, it shows as can't find a valid checkpoint. Could someone please post some code snippet on how to use resume from checkpoint for PEFT models?

As shared in the snippet above, to make `resume_from_checkpoint` work as expected, it assumes that you have previously trained your model using trainer that saves artifacts under `{output_dir}/checkpoint-{i}`, I have "faked" that in the example by manually saving a model in a folder called `{output_dir}/checkpoint-1`. Therefore you need to make sure the model weights lives under that folder.

> This doesn't seem like older resume from checkpoint that has it for pytorch models. Inside trainer.train we need to pass resume from checkpoint parameters as our last checkpoint path. while passing the path, it shows as can't find a valid checkpoint. Could someone please post some code snippet on how to use resume from checkpoint for PEFT models?

Hi @younesbelkada can you look at my issue with the code and please address it?


Yes @adityaaryan77 , sure, please have a look at my comment on the PEFT issue and discuss there

Hi @younesbelkada

> ```python
> ```python
> trainer.train(resume_from_checkpoint=True)
> ```
> 
> 
>     
>       
>     
> 
>       
>     
> 
>     
>   
> ```

Thank you for your response.

Please look at below code snippet:

# -*- coding: utf-8 -*-
"""Untitled345.ipynb

Automatically generated by Colaboratory.

Original file is located at
    https://colab.research.google.com/drive/1SgzMXUUDK1wDH0M0yQPfWmeNAKyy7EFs
"""

! pip install datasets transformers peft evaluate

!git clone https://github.com/llohann-speranca/transformers.git -b fix-resume-checkpoint-for-peftmodel

!cp -r /content/transformers /usr/local/lib/python3.10/dist-packages/transformers

import transformers
import numpy as np
GLUE_TASKS = ["cola", "mnli", "mnli-mm", "mrpc", "qnli", "qqp", "rte", "sst2", "stsb", "wnli"]
task = "cola"
model_checkpoint = "bert-large-uncased"
batch_size = 16
from datasets import load_dataset, load_metric
actual_task = "mnli" if task == "mnli-mm" else task
dataset = load_dataset("glue", actual_task)
metric = load_metric('glue', actual_task)
from transformers import AutoTokenizer
tokenizer = AutoTokenizer.from_pretrained(model_checkpoint, use_fast=True)
task_to_keys = {
    "cola": ("sentence", None),
    "mnli": ("premise", "hypothesis"),
    "mnli-mm": ("premise", "hypothesis"),
    "mrpc": ("sentence1", "sentence2"),
    "qnli": ("question", "sentence"),
    "qqp": ("question1", "question2"),
    "rte": ("sentence1", "sentence2"),
    "sst2": ("sentence", None),
    "stsb": ("sentence1", "sentence2"),
    "wnli": ("sentence1", "sentence2"),
}
sentence1_key, sentence2_key = task_to_keys[task]
if sentence2_key is None:
    print(f"Sentence: {dataset['train'][0][sentence1_key]}")
else:
    print(f"Sentence 1: {dataset['train'][0][sentence1_key]}")
    print(f"Sentence 2: {dataset['train'][0][sentence2_key]}")
def preprocess_function(examples):
    if sentence2_key is None:
        return tokenizer(examples[sentence1_key], truncation=True)
    return tokenizer(examples[sentence1_key], examples[sentence2_key], truncation=True)
encoded_dataset = dataset.map(preprocess_function, batched=True)
from transformers import AutoModelForSequenceClassification, TrainingArguments, Trainer
from peft import (
    get_peft_config,
    get_peft_model,
    get_peft_model_state_dict,
    set_peft_model_state_dict,
    LoraConfig,
    PeftType,
    PrefixTuningConfig,
    PromptEncoderConfig,
)
peft_type = PeftType.LORA
device = "cuda"
peft_config = LoraConfig(task_type="SEQ_CLS", inference_mode=False, r=8, lora_alpha=16, lora_dropout=0.1)
lr = 3e-4

num_labels = 3 if task.startswith("mnli") else 1 if task=="stsb" else 2
model = AutoModelForSequenceClassification.from_pretrained(model_checkpoint, num_labels=num_labels)
model = get_peft_model(model, peft_config)
model.print_trainable_parameters()
model
metric_name = "pearson" if task == "stsb" else "matthews_correlation" if task == "cola" else "accuracy"
model_name = model_checkpoint.split("/")[-1]

args = TrainingArguments(
    f"{model_name}-finetuned1-{task}",
    evaluation_strategy = "epoch",
    save_strategy = "epoch",
    learning_rate=2e-5,
    per_device_train_batch_size=batch_size,
    per_device_eval_batch_size=batch_size,
    num_train_epochs=2,
    weight_decay=0.01,
    # load_best_model_at_end=True,
    metric_for_best_model=metric_name,
    # push_to_hub=True,
)
from transformers import Seq2SeqTrainer, TrainerCallback, TrainingArguments, TrainerState, TrainerControl
from transformers.trainer_utils import PREFIX_CHECKPOINT_DIR
import os

class SavePeftModelCallback(TrainerCallback):
    def on_save(
        self,
        args: TrainingArguments,
        state: TrainerState,
        control: TrainerControl,
        **kwargs,
    ):
        checkpoint_folder = os.path.join(args.output_dir, f"{PREFIX_CHECKPOINT_DIR}-{state.global_step}")

        peft_model_path = os.path.join(checkpoint_folder, "adapter_model")
        kwargs["model"].save_pretrained(peft_model_path)

        pytorch_model_path = os.path.join(checkpoint_folder, "pytorch_model.bin")
        if os.path.exists(pytorch_model_path):
            os.remove(pytorch_model_path)
        return control
def compute_metrics(eval_pred):
    predictions, labels = eval_pred
    if task != "stsb":
        predictions = np.argmax(predictions, axis=1)
    else:
        predictions = predictions[:, 0]
    return metric.compute(predictions=predictions, references=labels)
validation_key = "validation_mismatched" if task == "mnli-mm" else "validation_matched" if task == "mnli" else "validation"
trainer = Trainer(
    model,
    args,
    train_dataset=encoded_dataset["train"],
    eval_dataset=encoded_dataset[validation_key],
    tokenizer=tokenizer,
    compute_metrics=compute_metrics,
    callbacks=[SavePeftModelCallback],
)
trainer.train()

trainer.save_model("/content/bert-large-uncased-finetuned1-cola/checkpoint-1")

trainer.train(resume_from_checkpoint='/content/bert-large-uncased-finetuned1-cola/checkpoint-1070/adapter_model')



Inside the resume from checkpoint i have tried with below options
1) resume_from_checkpoint = True
2) resume_from_checkpoint  = (Last checkpoint path)
3) resume_from_checkpoint  = (trainer.saved model path)

Everywhere I'm getting the same message of Can't find a valid checkpoint at <Model saved path>.
At the same time, I'm able to continue my resume from checkpoint in native pytorch code.


> Hi @llohann-speranca Again thanks for your great work on this, I think this seems a rather important fix that might unlock a lot of users, if that's ok for you, I can quickly take over the PR and address the last comment so that we can merge the PR. What do you think ?

Hi @younesbelkada. Sure! I have been very busy and have still to learn how to deal with PRs. Sorry about that.

@llohann-speranca thanks! 
@techthiyanes it seems you are using the API the wrong way. `resume_from_checkpoint` will try to retrieve the latest checkpoint from the output directory of the trainer. Therefore make sure you have correct `checkpoints-{i}` folders inside `f"{model_name}-finetuned1-{task}"` in your case and use `resume_from_checkpoint=True`

> @llohann-speranca thanks! @techthiyanes it seems you are using the API the wrong way. `resume_from_checkpoint` will try to retrieve the latest checkpoint from the output directory of the trainer. Therefore make sure you have correct `checkpoints-{i}` folders inside `f"{model_name}-finetuned1-{task}"` in your case and use `resume_from_checkpoint=True`

Thanks a lot on your response.
By default while passing resume from checkpoint then API automatically consumes the recent checkpoint. But this is something not working as expected for PEFT models than torch models.  As mentioned, I have pointed out the correct checkpoint and the same folder resides inside alone.

If you don't mind, could you please try executing the any of huggingface example code inserting PEFT with the trainer & resume from checkpoint? Then you might be able to replicate. 

@techthiyanes I think it works as expected with this PR, as explained in https://github.com/huggingface/transformers/pull/24274#pullrequestreview-1479614483 I have tried the attached snippet that was not working before the PR as mentioned and this PR properly fixes it by loading the checkpoint. If you want you can try to replicate using a smaller example (for example imdb as attached) and let me know if you still face an issue by opening a new ticket

> @techthiyanes I think it works as expected with this PR, as explained in [#24274 (review)](https://github.com/huggingface/transformers/pull/24274#pullrequestreview-1479614483) I have tried the attached snippet that was not working before the PR as mentioned and this PR properly fixes it by loading the checkpoint. If you want you can try to replicate using a smaller example (for example imdb as attached) and let me know if you still face an issue by opening a new ticket

Sure..Thanks a lot.. Let me try above snippet for classification models then let you know.

> @techthiyanes I think it works as expected with this PR, as explained in [#24274 (review)](https://github.com/huggingface/transformers/pull/24274#pullrequestreview-1479614483) I have tried the attached snippet that was not working before the PR as mentioned and this PR properly fixes it by loading the checkpoint. If you want you can try to replicate using a smaller example (for example imdb as attached) and let me know if you still face an issue by opening a new ticket

Still able to replicate the issue.  Raised a separate issue on the same.
https://github.com/huggingface/transformers/issues/24354

Hi, nice job! Does this new feature available if I `pip install` the latest peft/transformers packages? Or should I install from source?

Thank you so much @llohann-speranca and @younesbelkada  for adding this 🤗! @beyondguo, please install from source as this isn't yet part of the release.

Thanks for iterating! Will we perform inference in the same manner? Specifically, `peft` requires us to load `PeftConfig` via `adapter_config.json`. I saw that this is not saved with `trainer.save_model()`. Will we need to add `model.save_pretrained()` to use for inference?

I am trying to resume in this way:

```python
from peft import PeftModel, PeftConfig
from transformers import AutoModelForSeq2SeqLM

config = PeftConfig.from_pretrained("huggingface_path_TO_MY_PREVIOUSLY_TRAINED_LORA")
model = AutoModelForSeq2SeqLM.from_pretrained("google/flan-t5-large") # underlying model
model.enable_input_require_grads() # to make training possible
model = PeftModel.from_pretrained(model, "huggingface_path_TO_MY_PREVIOUSLY_TRAINED_LORA")

model.print_trainable_parameters()
# trainable params: 0 || all params: 792,587,264 || trainable%: 0.0
```

Then the usual code:
```python
from transformers import DataCollatorForSeq2Seq

# we want to ignore tokenizer pad token in the loss
label_pad_token_id = -100
# Data collator
data_collator = DataCollatorForSeq2Seq(
    tokenizer,
    model=model,
    label_pad_token_id=label_pad_token_id,
)

from transformers import Seq2SeqTrainer, Seq2SeqTrainingArguments

output_dir="./t5-large-r32-lora-JOIN-FIX-RESUME"
#batch_size = 8

# Define training args
training_args = Seq2SeqTrainingArguments(
    output_dir=output_dir,
    auto_find_batch_size=True,
    save_strategy="steps",
    save_steps=500,
    gradient_accumulation_steps=4,
    learning_rate=1e-3, # higher learning rate
    weight_decay=0.01,
    num_train_epochs=2, 
    logging_dir=f"{output_dir}/logs",
    logging_strategy="steps",
    logging_steps=100,
    push_to_hub=True)

# Create Trainer instance
trainer = Seq2SeqTrainer(
    model=model,
    args=training_args,
    data_collator=data_collator,
    train_dataset=TRAINING,
)
model.config.use_cache = False  # silence the warnings. Please re-enable for inference!

trainer.train()
```

> ***The loss starts where I paused the training***. Which is so low so that I can tell that it has started training the old model, but I am not sure whether I am doing it right.

Will @younesbelkada you please make me sure if whatever I am doing is right?
Thank you.



> I am trying to resume in this way:
> 
> ```python
> from peft import PeftModel, PeftConfig
> from transformers import AutoModelForSeq2SeqLM
> 
> config = PeftConfig.from_pretrained("huggingface_path_TO_MY_PREVIOUSLY_TRAINED_LORA")
> model = AutoModelForSeq2SeqLM.from_pretrained("google/flan-t5-large") # underlying model
> model.enable_input_require_grads() # to make training possible
> model = PeftModel.from_pretrained(model, "huggingface_path_TO_MY_PREVIOUSLY_TRAINED_LORA")
> 
> model.print_trainable_parameters()
> # trainable params: 0 || all params: 792,587,264 || trainable%: 0.0
> ```
> 
> Then the usual code:
> 
> ```python
> from transformers import DataCollatorForSeq2Seq
> 
> # we want to ignore tokenizer pad token in the loss
> label_pad_token_id = -100
> # Data collator
> data_collator = DataCollatorForSeq2Seq(
>     tokenizer,
>     model=model,
>     label_pad_token_id=label_pad_token_id,
> )
> 
> from transformers import Seq2SeqTrainer, Seq2SeqTrainingArguments
> 
> output_dir="./t5-large-r32-lora-JOIN-FIX-RESUME"
> #batch_size = 8
> 
> # Define training args
> training_args = Seq2SeqTrainingArguments(
>     output_dir=output_dir,
>     auto_find_batch_size=True,
>     save_strategy="steps",
>     save_steps=500,
>     gradient_accumulation_steps=4,
>     learning_rate=1e-3, # higher learning rate
>     weight_decay=0.01,
>     num_train_epochs=2, 
>     logging_dir=f"{output_dir}/logs",
>     logging_strategy="steps",
>     logging_steps=100,
>     push_to_hub=True)
> 
> # Create Trainer instance
> trainer = Seq2SeqTrainer(
>     model=model,
>     args=training_args,
>     data_collator=data_collator,
>     train_dataset=TRAINING,
> )
> model.config.use_cache = False  # silence the warnings. Please re-enable for inference!
> 
> trainer.train()
> ```
> 
> > _**The loss starts where I paused the training**_. Which is so low so that I can tell that it has started training the old model, but I am not sure whether I am doing it right.
> 
> Will @younesbelkada you please make me sure if whatever I am doing is right? Thank you.

I've encountered the same issue. Have you resolved it? @AayushSameerShah 

@XM-Dong Nah... it seems like LoRA needs some "special script" :(

Hi, I'm also wondering how can I get these changes? Are they in a new transformers version or PEFT version?

my current versions are:
```
transformers==4.30.1
peft==0.4.0
```