Unable to Resume Training from LoRA Checkpoints When Using FSDP
Author: fabianlim
Pull request: no

### System Info

transformers==4.35.2
accelerate==0.23.0
peft==0.5.0

`accelerate.yaml`

```yaml
compute_environment: LOCAL_MACHINE
distributed_type: FSDP
downcast_bf16: 'no'
fsdp_config:
  fsdp_auto_wrap_policy: TRANSFORMER_BASED_WRAP
  fsdp_backward_prefetch_policy: BACKWARD_PRE
  fsdp_forward_prefetch: true
  fsdp_offload_params: false
  fsdp_sharding_strategy: 1
  fsdp_state_dict_type: FULL_STATE_DICT
  fsdp_sync_module_states: true
  fsdp_transformer_layer_cls_to_wrap: "BertLayer"
machine_rank: 0
main_training_function: main
mixed_precision: 'no'
num_machines: 1
num_processes: 2
rdzv_backend: static
same_network: true
tpu_env: []
tpu_use_cluster: false
tpu_use_sudo: false
use_cpu: false
```

### Who can help?

@pacman100 Following the recommendation https://huggingface.co/docs/trl/v0.7.4/en/sft_trainer#training-adapters to install a `PeftSavingCallback` to ensure that `adapter.bin` is saved. This will be the case when using `FSDP` since it is not a `PretrainedModel`, in which only the `state_dict` will be saved.

The recommendation above works great for saving the checkpoint, but does not work when resuming the checkpoint. This is because `model_wrapped` is neither a `PretrainedModel` nora  `PEFTModel`, and the `if-else` conditions in `Trainer._load_from_checkpoint` will go all the way to `load_sharded_checkpoint`. This results in the following error:

```shell
Traceback (most recent call last):
  File "/dccstor/flim-ai4it/AI4IT/tofafm/src/scripts/lora_fsdp_bug_demo.py", line 183, in <module>
    main()
  File "/dccstor/flim-ai4it/AI4IT/tofafm/src/scripts/lora_fsdp_bug_demo.py", line 143, in main
    trainer.train(resume_from_checkpoint=True)
  File "/u/flim/miniconda3/envs/tofafm-rewrite2/lib/python3.10/site-packages/transformers/trainer.py", line 1555, in train
    return inner_training_loop(
  File "/u/flim/miniconda3/envs/tofafm-rewrite2/lib/python3.10/site-packages/transformers/trainer.py", line 1712, in _inner_training_loop
    self._load_from_checkpoint(resume_from_checkpoint, self.model_wrapped)
  File "/u/flim/miniconda3/envs/tofafm-rewrite2/lib/python3.10/site-packages/transformers/trainer.py", line 2132, in _load_from_checkpoint
    load_result = load_sharded_checkpoint(
  File "/u/flim/miniconda3/envs/tofafm-rewrite2/lib/python3.10/site-packages/transformers/modeling_utils.py", line 411, in load_sharded_checkpoint
    raise ValueError(f"Can't find a checkpoint index ({' or '.join(filenames)}) in {folder}.")
```

The second issue with the recommendation, is that the FSDP optimizer sates are not saved in the `PeftSavingCallback`, so it will not be a clean fix.

I was wondering if you may have any thoughts on this. A possible hacky solution will be to override `Trainer._load_from_checkpoint` and use `FSDP.summon_full_params` to unshard the LoRA weights, and then call `load_adapter`, but it doesnt sound very clean given that it will not resume the FSDP optimizer. 


### Information

- [ ] The official example scripts
- [X] My own modified scripts

### Tasks

- [X] An officially supported task in the `examples` folder (such as GLUE/SQuAD, ...)
- [ ] My own task or dataset (give details below)

### Reproduction

1. Run the below script with the above `accelerate.yaml` configurations. After at least 100 steps, when an `adapter.bin` checkpoint has been populated, stop.
2. Rerun the script but now enabling `trainer.train(resume_from_checkpoint=True) `.

```python
   
from datasets import load_dataset
from transformers import AutoTokenizer, AutoModelForSequenceClassification, DataCollatorWithPadding
from transformers import TrainingArguments, Trainer, TrainerCallback
import torch, os

def main(
    model_name: str = 'textattack/bert-base-uncased-SST-2',
):

    # we set the max sequence length here
    tokenizer = AutoTokenizer.from_pretrained(
        model_name, model_max_length=512,
    )

    # load and tokenize a verys mall dataset
    raw_datasets = load_dataset('glue','sst2')

    # tokenization function
    def _tokenize_function(example, tokenizer):    
        return tokenizer(
            example['sentence'], truncation = True,
        )

    tokenized_datasets = raw_datasets.map(
        _tokenize_function, fn_kwargs = {'tokenizer': tokenizer}, batched=True
    )

    data_collator = DataCollatorWithPadding(
        tokenizer=tokenizer, return_tensors='pt'
    )

    model = AutoModelForSequenceClassification.from_pretrained(model_name)

    from peft import LoraConfig, get_peft_model
    model = get_peft_model(
        model, LoraConfig(
            r=8, lora_alpha=16,
            target_modules=['query', 'key', 'value'],
            task_type='SEQ_CLS'
        )
    )

    training_args = TrainingArguments(
        num_train_epochs = 1,
        output_dir = './results',
        per_device_train_batch_size = 8,
        per_device_eval_batch_size = 8,
        learning_rate = 2e-4,
        logging_steps = 50,
        save_strategy = 'steps',
        save_steps = 100,
        evaluation_strategy = 'steps',
        eval_steps = 100,
        save_total_limit = 2,
        metric_for_best_model = 'loss',
        greater_is_better = False,
        max_steps = 1000, # just make the demo quit after 1000 steps
        save_safetensors=False,
    )

    class PeftSavingCallback(TrainerCallback):
        def on_save(self, args, state, control, **kwargs):
            checkpoint_path = os.path.join(args.output_dir, f"checkpoint-{state.global_step}")
            kwargs["model"].save_pretrained(checkpoint_path)

            if "pytorch_model.bin" in os.listdir(checkpoint_path):
                os.remove(os.path.join(checkpoint_path, "pytorch_model.bin"))

    trainer = Trainer(
        model=model,
        args=training_args,
        train_dataset=tokenized_datasets["train"],
        eval_dataset=tokenized_datasets["validation"],
        data_collator=data_collator,
        callbacks=[PeftSavingCallback()], 
    )

    import functools
    from accelerate import DistributedType

    if trainer.accelerator.distributed_type == DistributedType.FSDP:

        from torch.distributed.fsdp.wrap import lambda_auto_wrap_policy, _or_policy

        # inspired by
        # https://github.com/facebookresearch/llama-recipes/blob/main/src/llama_recipes/utils/fsdp_utils.py
        def lambda_policy_fn(module):
            if (
                len(list(module.named_children())) == 0
                and getattr(module, "weight", None) is not None
                and module.weight.requires_grad
            ):
                return True
            return False

        trainer.accelerator.state.fsdp_plugin.set_auto_wrap_policy(model)
        trainer.accelerator.state.fsdp_plugin.auto_wrap_policy = functools.partial(
            _or_policy, policies=[
                functools.partial(lambda_auto_wrap_policy, lambda_fn=lambda_policy_fn),
                trainer.accelerator.state.fsdp_plugin.auto_wrap_policy
            ])

    # checkpoints will be saved every 100 steps as `pytorch.bin`
    trainer.train()
    # trainer.train(resume_from_checkpoint=True) # activating this will throw the error
```

### Expected behavior

1. `resume_from_checkpoint=True` will resume the PEFT checkpoint recorded by `PeftSavingCallback`.
3. [bonus]: the FSDP optimizer states can be resumed also.

Comments:
This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the [contributing guidelines](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md) are likely to be ignored.

Hello @fabianlim, I think the PR https://github.com/huggingface/transformers/pull/28297 should resolve this. 

@pacman100 yes I think so too, closing this issue.