Using Trainer with torch.compile() and use_orig_params=True produce model checkpoints that cannot be loaded 
Author: ikergarcia1996
Pull request: no

### System Info

- `transformers` version: 4.32.0.dev0
- Platform: Linux-4.18.0-477.10.1.el8_8.x86_64-x86_64-with-glibc2.28
- Python version: 3.9.7
- Huggingface_hub version: 0.14.1
- Safetensors version: 0.3.1
- Accelerate version: 0.21.0
- Accelerate config: 	- compute_environment: LOCAL_MACHINE
	- distributed_type: NO
	- mixed_precision: fp16
	- use_cpu: False
	- num_processes: 1
	- machine_rank: 0
	- num_machines: 1
	- rdzv_backend: static
	- same_network: False
	- main_training_function: main
	- downcast_bf16: False
	- tpu_use_cluster: False
	- tpu_use_sudo: False
- PyTorch version (GPU?): 2.1.0.dev20230523+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?

@sgugger

### 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 am using the Huggingface Trainer to fine-tune an LLaMA2 model with FSDP. I launch the script with the following command. I set `--fsdp_use_orig_params true` because without it I cannot get `torch.compile()` + FSDP to work (https://github.com/huggingface/transformers/pull/23481)

```bash
accelerate launch --num_processes=4 \
--use_fsdp \
--mixed_precision=bf16 \
--fsdp_auto_wrap_policy=TRANSFORMER_BASED_WRAP  \
--fsdp_transformer_layer_cls_to_wrap="LlamaDecoderLayer" \
--fsdp_sharding_strategy=1 \
--fsdp_state_dict_type=FULL_STATE_DICT \
--fsdp_use_orig_params true \
src/run.py ${CONFIGS_FOLDER}/LLaMa2_FSDP.yaml
```

In the trainer configuration, I set the following parameters that are relevant to the issue.

```yalm
torch_dtype: float32
gradient_checkpointing: true
fsdp: full_shard auto_wrap
bf16: true
fp16: false
torch_compile: true
```

The model is compiled correctly by PyTorch, the training is fast, and the loss looks good. However, when the training ends, I use the following line to save the model:

```
trainer.save_model()
```

The first problem is that the trainer seems to save two copies of the model: one is split into multiple parts, and the other one contains the same model in a single .bin file.

```bash
ls output_path

checkpoint-10322  config.json                       pytorch_model-00002-of-00003.bin  pytorch_model.bin.index.json  tokenizer.json
checkpoint-15483  generation_config.json            pytorch_model-00003-of-00003.bin  special_tokens_map.json       tokenizer.model
checkpoint-5161   pytorch_model-00001-of-00003.bin  pytorch_model.bin                 tokenizer_config.json         training_args.bin
```

The second problem, which is related to torch.compile(), is that the model weights are saved with the _orig_mod prefix.
```json
{
  "metadata": {
    "total_size": 26953670656
  },
  "weight_map": {
    "_orig_mod.lm_head.weight": "pytorch_model-00003-of-00003.bin",
    "_orig_mod.model.embed_tokens.weight": "pytorch_model-00001-of-00003.bin",
    "_orig_mod.model.layers.0.input_layernorm.weight": "pytorch_model-00001-of-00003.bin",
    "_orig_mod.model.layers.0.mlp.down_proj.weight": "pytorch_model-00001-of-00003.bin",
    "_orig_mod.model.layers.0.mlp.gate_proj.weight": "pytorch_model-00001-of-00003.bin",
    "..."
}
}
```

So, when I try to load the model for inference using model: PreTrainedModel = AutoModelForCausalLM.from_pretrained(), I get a huge warning that says all the weights in the .bin file are not used, and all the LLaMA2 weights have been randomly initialized.

### Expected behavior

The trainer should save the uncompiled model or correctly handle the `_orig_mod` prefixes. 

Comments:
I think this comes from a bas interaction between `torch.compile` and FSDP, cc @pacman100 

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.

I am encountering the same issue, where after saving a model that has been compiled using `torch.compile`, `safetensors.load_model` throws: 
```
RuntimeError: Error(s) in loading state_dict for DummyModel:
    Missing key(s) in state_dict: "module.0.bias", "module.0.weight", "module.2.bias", "module.2.weight"
    Unexpected key(s) in state_dict: "_orig_mod.module.0.bias", "_orig_mod.module.0.weight", "_orig_mod.module.2.bias", "_orig_mod.module.2.weight"
```
In this case, the model has a `nn.Sequential` called `module`. Can draft a quick repro if you wish. 

A fix I found is to unwrap the model, but this only works if you know the module names a priori: 
```py
# https://github.com/huggingface/transformers/blob/main/src/transformers/modeling_utils.py#L4788C1-L4799C21
def unwrap_model(model: nn.Module) -> nn.Module:
    """
    Recursively unwraps a model from potential containers (as used in distributed training).

    Args:
        model (`torch.nn.Module`): The model to unwrap.
    """
    # since there could be multiple levels of wrapping, unwrap recursively
    if hasattr(model, "module"):
        return unwrap_model(model.module)
    else:
        return model

# ...

safetensors.save_model(unwrap_model(model))
safetensors.load_model(unwrap_model(model))
```


Any planned fix for this? @sgugger 

Hi @peacefulotter could you open a new issue, detailing the problem, the running environment and linking to this issue as this seems to be related to safetensors?