Loading a model fails if it has been compiled with torch.compile
Author: peacefulotter
Pull request: no

### System Info

- `transformers` version: 4.36.1
- Platform: Linux-6.2.0-37-generic-x86_64-with-glibc2.35
- Python version: 3.10.12
- Huggingface_hub version: 0.19.4
- Safetensors version: 0.4.1
- Accelerate version: not installed
- Accelerate config: not found
- PyTorch version (GPU?): 2.1.0+cu118 (False)
- 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?: NO
- Using distributed or parallel set-up in script?: NO

triton: 2.1.0
Ubuntu 22.04 (jammy, LTS)

### Who can help?

@muellerzr @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

Any class that inherits `nn.Module`
```py
import torch
import torch.nn as nn

class MyModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.module = nn.Sequential(
            nn.Linear(4, 2),
            nn.ReLU(),
        )

    def forward(self, x):
        return self.module(x)

# Instantiate the model, and compile it
model = MyModel()
model = torch.compile(model)

# ... train or do whatever

# save it (using safetensors in my case)
import safetensors.torch as st
st.save_model(model,  "model.safetensors") 
# And load the weights, using safetensor as well
# The following throws a RuntimeError (see below)
st.load_model(model, "model.safetensors")
```

### Expected behavior

I am encountering a similar issue as in https://github.com/huggingface/transformers/issues/25205, 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", ...
    Unexpected key(s) in state_dict: "_orig_mod.module.0.bias", "_orig_mod.module.0.weight", ...
```
In this case, the model has a `nn.Sequential` called `module`. As one can see, loading the weights changes the layer names by adding `_orig_mod` at the front. 

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

# ...

st.save_model(unwrap_model(model), "model.safetensors")
st.load_model(unwrap_model(model), "model.safetensors")
```

Comments:
This would be better posted in the safetensors repo rather than transformers I believe? https://github.com/huggingface/safetensors

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.

For anyone stumbling upon this: if you save the weights of a *compiled* model then you also need to load the weights inside that *compiled* model. 

The repro above is actually supposed to work..