Model state dict keys mismatch when set `load_best_model_at_end`=True using LoRA and deepspeed
Author: liuchengyuan123
Pull request: no

### System Info

I am running with transformers.trainer, fine-tuning the Qwen1.5-1.8B-Chat model with LoRA. The lora targets are [q_proj, v_proj]. I also set the lm_head and embedding layer as trainable (although it is not the key to the error I guess).

However, when I try to load the best checkpoint from disk, I get the following error:

```bash
File "x/src/llmtuner/train/sft/workflow.py", line 134, in run_sft
    train_result = trainer.train(resume_from_checkpoint=training_args.resume_from_checkpoint)
  File "x/site-packages/transformers/trainer.py", line 1539, in train
    return inner_training_loop(
  File "x/site-packages/transformers/trainer.py", line 1972, in _inner_training_loop
    self._load_best_model()
  File "x/site-packages/transformers/trainer.py", line 2168, in _load_best_model
    deepspeed_load_checkpoint(self.model_wrapped, self.state.best_model_checkpoint)
  File "x/site-packages/transformers/integrations/deepspeed.py", line 402, in deepspeed_load_checkpoint
    load_path, _ = deepspeed_engine.load_checkpoint(
  File "/x/site-packages/deepspeed/runtime/engine.py", line 2697, in load_checkpoint
    load_path, client_states = self._load_checkpoint(load_dir,
  File "x/site-packages/deepspeed/runtime/engine.py", line 2762, in _load_checkpoint
    self.load_module_state_dict(checkpoint=checkpoint,
  File "x/site-packages/deepspeed/runtime/engine.py", line 2560, in load_module_state_dict
    self.module.load_state_dict(
  File "x/site-packages/torch/nn/modules/module.py", line 2152, in load_state_dict
    raise RuntimeError('Error(s) in loading state_dict for {}:\n\t{}'.format(
RuntimeError: Error(s) in loading state_dict for PeftModelForCausalLM:
	Missing key(s) in state_dict: "base_model.model.model.layers.0.self_attn.q_proj.base_layer.weight", "base_model.model.model.layers.0.self_attn.q_proj.base_layer.bias", "base_model.model.model.layers.0.self_attn.k_proj.weight", "base_model.model.model.layers.0.self_attn.k_proj.bias", "base_model.model.model.layers.0.self_attn.v_proj.base_layer.weight", "base_model.model.model.layers.0.self_attn.v_proj.base_layer.bias", "base_model.model.model.layers.0.self_attn.o_proj.weight", "base_model.model.model.layers.0.mlp.gate_proj.weight", "base_model.model.model.layers.0.mlp.up_proj.weight", "base_model.model.model.layers.0.mlp.down_proj.weight", "base_model.model.model.layers.0.input_layernorm.weight", "base_model.model.model.layers.0.post_attention_layernorm.weight", "base_model.model.model.layers.1.self_attn.q_proj.base_layer.weight", "base_model.model.model.layers.1.self_attn.q_proj.base_layer.bias", "base_model.model.model.layers.1.self_attn.k_proj.weight", "base_model.model.model.layers.1.self_attn.k_proj.bias...
```

While I succeeded to load the adapter state dict using `PEFTModel.load_adapter`, I assume it may be a problem of intergrating deepspeed with PEFT?

### Who can help?

_No response_

### Information

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

### Tasks

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

### Reproduction

environments:

transformers==4.37.2
accelerate==0.26.1
peft==0.8.2
gradio==3.50.2
deepspeed==0.11.1
(I have tried both deepspeed of 0.11 and 0.13, facing the same problem)

deepspeed_config.json:
```json
{
    "fp16": {
        "enabled": "auto",
        "loss_scale": 0,
        "loss_scale_window": 1000,
        "initial_scale_power": 16,
        "hysteresis": 2,
        "min_loss_scale": 1
    },
    "zero_optimization": {
        "stage": 2,
        "allgather_partitions": true,
        "allgather_bucket_size": 5e8,
        "reduce_scatter": true,
        "reduce_bucket_size": 5e8,
        "overlap_comm": false,
        "contiguous_gradients": true
    },
    "gradient_accumulation_steps": "auto",
    "gradient_clipping": "auto",
    "train_batch_size": "auto",
    "train_micro_batch_size_per_gpu": "auto",
    "steps_per_print": 1000
}
```

training_arguments:
```bash
--load_best_model_at_end true
```

### Expected behavior

All state dict keys are matched.

Comments:
I just encountered the same issue. Previously, I was using transformers==4.34.0 and deepspeed==0.12.6 without any problems. Recently, I upgraded the packages to deepspeed==0.13.4 and transformers==4.37.2, and then I started experiencing the same problem as you. Based on your information, I am considering that it might be an issue with transformers, and I am trying to revert to an earlier version of transformers to solve it.

After trying, rolling back the Transformers version to 4.36.2 seems to have solved this problem

cc @younesbelkada and @pacman100 

i think it should be now supported if you use latest transformers / accelerate / peft, see: https://huggingface.co/docs/peft/accelerate/deepspeed#use-peft-and-deepspeed-with-zero3-for-finetuning-large-models-on-multiple-machines-and-multiple-nodes

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.