low_cpu_mem_usage raises KeyError with modified GPT2 model
Author: Wenhan-Tan
Pull request: no

### System Info

```
- `transformers` version: 4.25.1
- Platform: Linux-5.4.0-135-generic-x86_64-with-glibc2.10
- Python version: 3.8.13
- Huggingface_hub version: 0.11.1
- PyTorch version (GPU?): 1.13.0+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?: Not yet
- Using distributed or parallel set-up in script?: Not yet
```

### Who can help?

_No response_

### 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 trying to test GPT2 models with different layer numbers, head numbers, and head sizes. The following code works with no errors. And the model is loaded successfully into the CPU with random weights, which is expected.
```
import torch
from transformers import AutoModelForCausalLM, AutoConfig

if __name__ == "__main__":
    model_id = "gpt2"
    model_config = AutoConfig.from_pretrained(pretrained_model_name_or_path=model_id)

    model_config.n_layer = 48
    model_config.n_head = 25
    model_config.n_embd = 1600
    model = AutoModelForCausalLM.from_pretrained(pretrained_model_name_or_path=model_id,
                                                 config=model_config,
                                                 ignore_mismatched_sizes=True,
                                                 torch_dtype=torch.float16)
```
However, when I set the flag `low_cpu_mem_usage=True` in `from_pretrained()` like this:
```
import torch
from transformers import AutoModelForCausalLM, AutoConfig

if __name__ == "__main__":
    model_id = "gpt2"
    model_config = AutoConfig.from_pretrained(pretrained_model_name_or_path=model_id)

    model_config.n_layer = 48
    model_config.n_head = 25
    model_config.n_embd = 1600
    model = AutoModelForCausalLM.from_pretrained(pretrained_model_name_or_path=model_id,
                                                 config=model_config,
                                                 ignore_mismatched_sizes=True,
                                                 torch_dtype=torch.float16,
                                                 low_cpu_mem_usage=True)
```
I get below errors:
```
/opt/conda/lib/python3.8/site-packages/scipy/__init__.py:138: UserWarning: A NumPy version >=1.16.5 and <1.23.0 is required for this version of SciPy (detected version 1.23.5)
  warnings.warn(f"A NumPy version >={np_minversion} and <{np_maxversion} is required for this version of "
Traceback (most recent call last):
  File "tmp.py", line 11, in <module>
    model = AutoModelForCausalLM.from_pretrained(pretrained_model_name_or_path=model_id,
  File "/home/wenhant/.local/lib/python3.8/site-packages/transformers/models/auto/auto_factory.py", line 463, in from_pretrained
    return model_class.from_pretrained(
  File "/home/wenhant/.local/lib/python3.8/site-packages/transformers/modeling_utils.py", line 2379, in from_pretrained
    ) = cls._load_pretrained_model(
  File "/home/wenhant/.local/lib/python3.8/site-packages/transformers/modeling_utils.py", line 2512, in _load_pretrained_model
    param = model_state_dict[key]
KeyError: 'h.45.attn.c_proj.bias'
```

### Expected behavior

I expect my code to run with no errors doesn't matter if I set `low_cpu_mem_usage` to `True` or `False`.

Comments:
Hi, @Wenhan-Tan 
I have made a PR regarding this issue, you can checkout the branch `fix_low_cpu_mem_usage` from my repository ([here](https://github.com/susnato/transformers/tree/fix_low_cpu_mem_usage)) and check if it solves your issue or not until the mods take any action on my PR or maybe merge it.

Thanks, 
susnato.

Hi @susnato ,
Thank you! Your PR solves the issue! But I get another one when I use DeepSpeed inference afterwards. Not sure if they're related. Code is below:
```
import torch
from transformers import AutoModelForCausalLM, AutoConfig
import deepspeed

if __name__ == "__main__":
    model_id = "gpt2"
    model_config = AutoConfig.from_pretrained(pretrained_model_name_or_path=model_id)

    model_config.n_layer = 48
    model_config.n_head = 25
    model_config.n_embd = 1600
    model = AutoModelForCausalLM.from_pretrained(pretrained_model_name_or_path=model_id,
                                                 config=model_config,
                                                 ignore_mismatched_sizes=True,
                                                 torch_dtype=torch.float16,
                                                 low_cpu_mem_usage=True)
    ds_config = {
        "tensor_parallel": {"tp_size": 1},
        "dtype": "fp16",
        "replace_with_kernel_inject": True,
        "replace_method": "auto",
    }
    ds_model = deepspeed.init_inference(model=model, config=ds_config)
```
I get errors below:
```
Traceback (most recent call last):
  File "tmp.py", line 23, in <module>
    ds_model = deepspeed.init_inference(model=model, config=ds_config)
  File "/home/wenhant/.local/lib/python3.8/site-packages/deepspeed/__init__.py", line 311, in init_inference
    engine = InferenceEngine(model, config=ds_inference_config)
  File "/home/wenhant/.local/lib/python3.8/site-packages/deepspeed/inference/engine.py", line 127, in __init__
    self.module.to(device)
  File "/home/wenhant/.local/lib/python3.8/site-packages/transformers/modeling_utils.py", line 1682, in to
    return super().to(*args, **kwargs)
  File "/home/wenhant/.local/lib/python3.8/site-packages/torch/nn/modules/module.py", line 987, in to
    return self._apply(convert)
  File "/home/wenhant/.local/lib/python3.8/site-packages/torch/nn/modules/module.py", line 639, in _apply
    module._apply(fn)
  File "/home/wenhant/.local/lib/python3.8/site-packages/torch/nn/modules/module.py", line 639, in _apply
    module._apply(fn)
  File "/home/wenhant/.local/lib/python3.8/site-packages/torch/nn/modules/module.py", line 662, in _apply
    param_applied = fn(param)
  File "/home/wenhant/.local/lib/python3.8/site-packages/torch/nn/modules/module.py", line 985, in convert
    return t.to(device, dtype if t.is_floating_point() or t.is_complex() else None, non_blocking)
NotImplementedError: Cannot copy out of meta tensor; no data!
```
This error won't occur if I don't use the flag `low_cpu_mem_usage=True`.