Peft Model not resuming from Checkpoint
Author: llohann-speranca
Pull request: no

### System Info

Running from huggingface/transformers-pytorch-gpu docker image.

- `transformers` version: 4.31.0.dev0
- Platform: Linux-5.4.0-139-generic-x86_64-with-glibc2.29
- Python version: 3.8.10
- Huggingface_hub version: 0.15.1
- Safetensors version: 0.3.1
- PyTorch version (GPU?): 2.0.1+cu118 (True)
- Tensorflow version (GPU?): not installed (NA)
- Flax version (CPU?/GPU?/TPU?): not installed (NA)
- Jax version: not installed
- JaxLib version: not installed
- Using GPU in script?: Yes
- Using distributed or parallel set-up in script?: No

### Who can help?

@younesbelkada

Since PR #24073 the Trainer does not resume from checkpoint.

The Issue happens since `PeftModel.save_pretrained` saves only adapter's files, but the `model._load_from_checkpoint` method expects a full pytorch checkpoint.

I worked around that by subclassing the Trainer class. I am willing to submit a PR merging the `_load_from_peft_checkpoint` with the Hugging face Trainer. 



```python
class PeftTrainer(Trainer):
 
    def _load_from_peft_checkpoint(self, resume_from_checkpoint, model):
        adapter_weights_file = os.path.join(resume_from_checkpoint, ADAPTER_WEIGHTS_NAME)
        adapter_safe_weights_file = os.path.join(resume_from_checkpoint, ADAPTER_SAFE_WEIGHTS_NAME)

        if not any(
            os.path.isfile(f) for f in [adapter_weights_file, adapter_safe_weights_file]
        ):
            raise ValueError(f"Can't find a valid checkpoint at {resume_from_checkpoint}")

        logger.info(f"Loading model from {resume_from_checkpoint}.")
        # Load adapters following PR # 24096 
        if is_peft_available() and isinstance(model, PeftModel):
            # If train a model using PEFT & LoRA, assume that adapter have been saved properly.
            if hasattr(model, "active_adapter") and hasattr(model, "load_adapter"):
                if os.path.exists(resume_from_checkpoint) or os.path.exists(resume_from_checkpoint):
                    model.load_adapter(resume_from_checkpoint, model.active_adapter)
                    # Load_adapter has no return value present, modify it when appropriate.
                    from torch.nn.modules.module import _IncompatibleKeys

                    load_result = _IncompatibleKeys([], [])
                else:
                    logger.warning(
                        "The intermediate checkpoints of PEFT may not be saved correctly, "
                        f"using `TrainerCallback` to save {ADAPTER_WEIGHTS_NAME} in corresponding folders, "
                        "here are some examples https://github.com/huggingface/peft/issues/96"
                    )
            else:
                logger.warning("Could not load adapter model, make sure to have `peft>=0.3.0` installed")

    def _load_from_checkpoint(self, resume_from_checkpoint, model=None):

        if model is None:
            model = self.model_wrapped if is_sagemaker_mp_enabled() else self.model

        if is_peft_available() and isinstance(model, PeftModel):
            # Try to load adapters before trying to load a torch model
            try:
                return self._load_from_peft_checkpoint(resume_from_checkpoint, model=model)
            except:
                return super()._load_from_checkpoint(resume_from_checkpoint, model=model)
            # If it is not a PeftModel, use the original _load_from_checkpoint
        else:
            return super()._load_from_checkpoint(resume_from_checkpoint, model=model)
     
```

### Information

- [ ] The official example scripts
- [X] My own modified scripts

### Tasks

- [ ] An officially supported task in the `examples` folder (such as GLUE/SQuAD, ...)
- [X] My own task or dataset (give details below)

### Reproduction

```python

from transformers import (
    AutoModelForSequenceClassification,
    AutoTokenizer,
    TrainingArguments,
    Trainer,
    DataCollatorWithPadding,
)
from peft import get_peft_config, get_peft_model, PromptTuningInit, PromptTuningConfig, TaskType, PeftType, PromptEncoderConfig
import torch
import os
import evaluate
from datasets import Dataset


# P-tuning hyper-parameters

model_id = "microsoft/deberta-v3-base"

model = AutoModelForSequenceClassification.from_pretrained(model_id)
tokenizer = AutoTokenizer.from_pretrained(model_id)

# Load tokenized datesets
train_ds = test_ds = Dataset.from_dict({'input_ids':[[1,  2430,  429,  92207,  303,  331,  1789,  3495,  2344,  1300,  355,  268,  1131,  270,  310,  354,  3732,  388,
  2],[1, 1865, 843, 20060, 265, 483, 2196, 281, 411, 2784, 2]],
'labels':[0,0]})

PEFT_CONFIG ={"peft_type":"P_TUNING",
          "num_virtual_tokens": 30,
              "encoder_reparameterization_type": "MLP",
              "encoder_hidden_size": 128,
                  "num_attention_heads": 17,

             }


peft_config = PromptEncoderConfig(
    task_type="SEQ_CLS",
    **PEFT_CONFIG
)

model = get_peft_model(model, peft_config)

data_collator = DataCollatorWithPadding(tokenizer=tokenizer, padding=True, max_length=482,)


training_args = TrainingArguments(
    output_dir='p',
    per_device_train_batch_size=1,
    per_device_eval_batch_size=1,   
    num_train_epochs=1,
    load_best_model_at_end=False,
    save_strategy='epoch'
)



trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=train_ds,
    eval_dataset=test_ds,
    tokenizer=tokenizer,
    data_collator=data_collator,
    compute_metrics=evaluate.load('accuracy')
)

trainer.train()


training_args = TrainingArguments(
    output_dir='p',
    per_device_train_batch_size=1,
    per_device_eval_batch_size=1,   
    num_train_epochs=2,
    load_best_model_at_end=False,
    save_strategy='epoch'

)



trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=train_ds,
    eval_dataset=test_ds,
    tokenizer=tokenizer,
    data_collator=data_collator,
    compute_metrics=evaluate.load('accuracy')
)

trainer.train(resume_from_checkpoint=True)

```

Raises

```
ValueError                                Traceback (most recent call last)
Cell In[26], line 92
     70 training_args = TrainingArguments(
     71     output_dir='p',
     72     per_device_train_batch_size=1,
   (...)
     77 
     78 )
     82 trainer = Trainer(
     83     model=model,
     84     args=training_args,
   (...)
     89     compute_metrics=evaluate.load('accuracy')
     90 )
---> 92 trainer.train(resume_from_checkpoint=True)

File /transformers/src/transformers/trainer.py:1634, in Trainer.train(self, resume_from_checkpoint, trial, ignore_keys_for_eval, **kwargs)
   1631         raise ValueError(f"No valid checkpoint found in output directory ({args.output_dir})")
   1633 if resume_from_checkpoint is not None and not is_sagemaker_mp_enabled() and not self.is_deepspeed_enabled:
-> 1634     self._load_from_checkpoint(resume_from_checkpoint)
   1636 # If model was re-initialized, put it on the right device and update self.model_wrapped
   1637 if model_reloaded:

File /transformers/src/transformers/trainer.py:2119, in Trainer._load_from_checkpoint(self, resume_from_checkpoint, model)
   2114 safe_weights_index_file = os.path.join(resume_from_checkpoint, SAFE_WEIGHTS_INDEX_NAME)
   2116 if not any(
   2117     os.path.isfile(f) for f in [weights_file, safe_weights_file, weights_index_file, safe_weights_index_file]
   2118 ):
-> 2119     raise ValueError(f"Can't find a valid checkpoint at {resume_from_checkpoint}")
   2121 logger.info(f"Loading model from {resume_from_checkpoint}.")
   2123 if os.path.isfile(config_file):

ValueError: Can't find a valid checkpoint at p/checkpoint-2
```




### Expected behavior

The train should be resumed from Epoch 1 and proceeded up to Epoch 2.
 

Comments:
Hi @llohann-speranca 
Thanks for digging, flagging the issue and proposing a fix! 
Indeed didn't properly tried it with `resume_from_checkpoint`. Yes please could you open a PR and tag me there? 
Thanks and looking forward to the PR!

Hi Younes. Thank you for the reply. I will do it later today, after working
hours.

On Tue, Jun 13, 2023 at 6:05 PM Younes Belkada ***@***.***>
wrote:

> Hi @llohann-speranca <https://github.com/llohann-speranca>
> Thanks for digging, flagging the issue and proposing a fix!
> Indeed didn't properly tried it with resume_from_checkpoint. Yes please
> could you open a PR and tag me there?
> Thanks and looking forward to the PR!
>
> —
> Reply to this email directly, view it on GitHub
> <https://github.com/huggingface/transformers/issues/24252#issuecomment-1589610039>,
> or unsubscribe
> <https://github.com/notifications/unsubscribe-auth/AZFKQJWCDRDMFWVI6OYNZ23XLCFVLANCNFSM6AAAAAAZFCV3FM>
> .
> You are receiving this because you were mentioned.Message ID:
> ***@***.***>
>