Fail to resume training from Pytorch FSDP checkpoint.
Author: Aria-K-Alethia
Pull request: no

### System Info

- `transformers` version: 4.35.2  (I also tried the development version 4.36.0 but the same problem arised)
- Platform: Linux-5.15.0-1033-azure-x86_64-with-glibc2.31
- Python version: 3.10.13
- Huggingface_hub version: 0.19.4
- Safetensors version: 0.4.1
- Accelerate version: 0.25.0
- Accelerate config:    not found
- PyTorch version (GPU?): 2.1.1+rocm5.6 (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?: I'm not sure the meaning of this term but I use 4 GPU cards

### Who can help?

@muellerzr @pacman100

### 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

**Problem description**
There is no direct code sample to reproduce, so I will try my best to describe how the problem appeared.

1. I fine-tuned a LLM using `transformers.trainer` with pytorch FSDP arguments. The cmd is like:
```python
torchrun --nproc_per_node=4 --master_port=20001 train.py \
    --model_name_or_path meta-llama/Llama-2-7b-hf  \
    --bf16 True \
    --output_dir output \
    --num_train_epochs 5 \
    --per_device_train_batch_size 2 \
    --per_device_eval_batch_size 16 \
    --gradient_accumulation_steps 16 \
    --evaluation_strategy "no" \
    --save_strategy "steps" \
    --save_steps 100 \
    --save_total_limit 8 \
    --learning_rate 1e-4 \
    --weight_decay 0. \
    --warmup_steps 1024 \
    --lr_scheduler_type "cosine" \
    --fsdp "full_shard auto_wrap" \
    --fsdp_transformer_layer_cls_to_wrap 'LlamaDecoderLayer' \
    --tf32 False \
    --logging_steps 1 \
    --model_max_length 4096 \
```
2. The training was interrupted so I tried to resume it from a checkpoint by `trainer.train(resume_from_checkpoint)`, but this gave the following error message:
```python
Traceback (most recent call last):
  File "/home/user/lm_finetune/train.py", line 183, in <module>
    train()
  File "/home/user/lm_finetune/train.py", line 175, in train
    trainer.train(resume_from_checkpoint=ckpt_dir)
  File "/opt/conda/envs/alpaca/lib/python3.10/site-packages/transformers/trainer.py", line 1555, in train
Traceback (most recent call last):
      File "/home/user/lm_finetune/train.py", line 183, in <module>
return inner_training_loop(
  File "/opt/conda/envs/alpaca/lib/python3.10/site-packages/transformers/trainer.py", line 1712, in _inner_training_loop
    train()
  File "/home/user/lm_finetune/train.py", line 175, in train
    trainer.train(resume_from_checkpoint=ckpt_dir)
  File "/opt/conda/envs/alpaca/lib/python3.10/site-packages/transformers/trainer.py", line 1555, in train
    self._load_from_checkpoint(resume_from_checkpoint, self.model_wrapped)
  File "/opt/conda/envs/alpaca/lib/python3.10/site-packages/transformers/trainer.py", line 2064, in _load_from_checkpoint
    return inner_training_loop(
  File "/opt/conda/envs/alpaca/lib/python3.10/site-packages/transformers/trainer.py", line 1712, in _inner_training_loop
    raise ValueError(f"Can't find a valid checkpoint at {resume_from_checkpoint}")
ValueError: Can't find a valid checkpoint at /home/user/lm_finetune/output/checkpoint-900
    self._load_from_checkpoint(resume_from_checkpoint, self.model_wrapped)
  File "/opt/conda/envs/alpaca/lib/python3.10/site-packages/transformers/trainer.py", line 2064, in _load_from_checkpoint
    raise ValueError(f"Can't find a valid checkpoint at {resume_from_checkpoint}")
ValueError: Can't find a valid checkpoint at /home/user/lm_finetune/output/checkpoint-900
```

**Possible solutions I tried**

1. I  first tried to update `transformers` to the development version `4.36.0dev`, but the problem still exists
2. I checked the source code, looks like it requires the model ckpt to be a folder, but in my case, the FSDP ckpt is only a file named `pytorch_model_fsdp.bin` . I think this should be the key point but I don't know how to solve it.

### Expected behavior

The training should be resumed.

Comments:
Hello, please let us know if the above PR resolves the issue.

Thank you for your follow up.
I tried the new branch, but it gave another error:
```python
Traceback (most recent call last):
  File "/home/user/lm_finetune/train_tts.py", line 183, in <module>
    train()
  File "/home/user/lm_finetune/train_tts.py", line 175, in train
    trainer.train(resume_from_checkpoint=ckpt_dir)
  File "/opt/conda/envs/alpaca/lib/python3.10/site-packages/transformers/trainer.py", line 1533, in train
    return inner_training_loop(
  File "/opt/conda/envs/alpaca/lib/python3.10/site-packages/transformers/trainer.py", line 1691, in _inner_training_loop
    self._load_optimizer_and_scheduler(resume_from_checkpoint)
  File "/opt/conda/envs/alpaca/lib/python3.10/site-packages/transformers/trainer.py", line 2518, in _load_optimizer_and_scheduler
    load_fsdp_optimizer(
  File "/opt/conda/envs/alpaca/lib/python3.10/site-packages/accelerate/utils/fsdp_utils.py", line 190, in load_fsdp_optimizer
    optimizer.load_state_dict(flattened_osd)
  File "/opt/conda/envs/alpaca/lib/python3.10/site-packages/accelerate/optimizer.py", line 107, in load_state_dict
    self.optimizer.load_state_dict(state_dict)
  File "/opt/conda/envs/alpaca/lib/python3.10/site-packages/torch/_compile.py", line 24, in inner
    return torch._dynamo.disable(fn, recursive)(*args, **kwargs)
  File "/opt/conda/envs/alpaca/lib/python3.10/site-packages/torch/_dynamo/eval_frame.py", line 328, in _fn
    return fn(*args, **kwargs)
  File "/opt/conda/envs/alpaca/lib/python3.10/site-packages/torch/optim/optimizer.py", line 735, in load_state_dict
    raise ValueError("loaded state dict contains a parameter group "
ValueError: loaded state dict contains a parameter group that doesn't match the size of optimizer's group
```