FSDP cuda out of memory during checkpoint saving
### System Info

- `transformers` version: 4.30.0.dev0
- Platform: Linux-4.14.81.bm.22-amd64-x86_64-with-glibc2.24
- Python version: 3.9.16
- Huggingface_hub version: 0.14.1
- Safetensors version: not installed
- PyTorch version (GPU?): 2.0.1+cu117 (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?: yes


### Who can help?

@stas00 @pacman100 

### 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

I'm running https://github.com/tatsu-lab/stanford_alpaca/blob/main/train.py on 32GB GPUs to finetune LLaMA-7B with FSDP turned on. The training process goes well but I get cuda oom when saving the final checkpoint. The entire fp32 state dict cannot fit in my gpu memory. A possible solution is to offload state dict to cpu mentioned in https://github.com/pytorch/pytorch/issues/98823? Is there any better way to handle this?

```
/data00/home/lijiahao.plus/miniconda3/envs/mlir/lib/python3.9/site-packages/torch/distributed/fsdp/_state_dict_utils.py:312: UserWarning: Failed to clone() tensor with name _fsdp_wrapped_module.model.layers.31.mlp.gate_proj.weight on rank 2. This may mean that this state_dict entry could point to invalid memory regions after returning from state_dict() call if this parameter is managed by FSDP. Please check clone implementation of _fsdp_wrapped_module.model.layers.31.mlp.gate_proj.weight. Error: CUDA out of memory. Tried to allocate 172.00 MiB (GPU 2; 31.75 GiB total capacity; 29.54 GiB already allocated; 39.75 MiB free; 30.38 GiB reserved in total by PyTorch) If reserved memory is >> allocated memory try setting max_split_size_mb to avoid fragmentation.  See documentation for Memory Management and PYTORCH_CUDA_ALLOC_CONF
```


### Expected behavior

Not oom on saving checkpoints.
Not sure why you tagged myself or Sourab - we have worked on Deepspeed integration, and you're welcome to ask questions if you use that. I personally don't know anything about FSDP - Deepspeed works perfectly well and FSDP implements the same ZeRO protocol that Deepspeed innovated. <|||||>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.