Models trained using Deepspeed ZeRO stage 3 have corrupted model weight shape
### System Info

transformers version: 4.21.1 | 4.24.0
- Platform: Linux-5.15.0-52-generic-x86_64-with-glibc2.17
- Python version: 3.8.13
- Huggingface_hub version: 0.10.0
- PyTorch version (GPU?): 1.12.1+cu113 (True)
- Tensorflow version (GPU?): 2.10.0 (True)
- 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

- [X] The official example scripts
- [ ] 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

I am currently trying to use deepspeed to finetune a AutoModelForCausalLM model (facebook/opt1.3b) on a multi-GPU instance with ZeRO optimization with the unmodified `run_clm_no_trainer.py` script from the examples. When I use ZeRO stage 2 to train the model, the model weights can be loaded normally. However, when I try using ZeRO stage 3 with CPU offloads for the optimizer weights, the model training proceeds normally with loss values and metrics that make sense. But I get the follow error when I try loading the weights.
```
RuntimeError: Error(s) in loading state_dict for OPTForCausalLM:
        size mismatch for model.decoder.embed_tokens.weight: copying a param with shape torch.Size([0]) from checkpoint, the shape in current model is torch.Size([50272, 2560]).
        size mismatch for model.decoder.embed_positions.weight: copying a param with shape torch.Size([0]) from checkpoint, the shape in current model is torch.Size([2050, 2560]).
        size mismatch for model.decoder.layers.0.self_attn.k_proj.weight: copying a param with shape torch.Size([0]) from checkpoint, the shape in current model is torch.Size([2560, 2560]).
        size mismatch for model.decoder.layers.0.self_attn.v_proj.weight: copying a param with shape torch.Size([0]) from checkpoint, the shape in current model is torch.Size([2560, 2560]).
        size mismatch for model.decoder.layers.0.self_attn.q_proj.weight: copying a param with shape torch.Size([0]) from checkpoint, the shape in current model is torch.Size([2560, 2560]).
        size mismatch for model.decoder.layers.0.self_attn.out_proj.weight: copying a param with shape torch.Size([0]) from checkpoint, the shape in current model is torch.Size([2560, 2560]).
...
        size mismatch for model.decoder.layers.31.fc1.weight: copying a param with shape torch.Size([0]) from checkpoint, the shape in current model is torch.Size([10240, 2560]).
        size mismatch for model.decoder.layers.31.fc2.weight: copying a param with shape torch.Size([0]) from checkpoint, the shape in current model is torch.Size([2560, 10240]).
        size mismatch for lm_head.weight: copying a param with shape torch.Size([0]) from checkpoint, the shape in current model is torch.Size([50272, 2560]).
        You may consider adding `ignore_mismatched_sizes=True` in the model `from_pretrained` method.
```
This is very strange as the `torch.Size([0])` error seems to be pervasive across all layers of the model, suggesting the weights are just empty and uninitialized. This is just speculation as the documentation does not seem to address the specifics of training with different ZeRO stages. I have tried the loading the model manually using `AutoModelForCausalLM.from_pretrained('./model_dir')` where `model_dir` is where the weights were saved after training, yet the same error is still thrown. I am not sure if this is a bug or using ZeRO stage 3 is currently unsupported. Any help would be much appreciated. 

### Expected behavior

Models trained using ZeRO stage 3 should load correctly.
Full error message with all the incorrect layers below:
```
RuntimeError: Error(s) in loading state_dict for OPTForCausalLM:
        size mismatch for model.decoder.embed_tokens.weight: copying a param with shape torch.Size([0]) from checkpoint, the shape in current model is torch.Size([50272, 2560]).
        size mismatch for model.decoder.embed_positions.weight: copying a param with shape torch.Size([0]) from checkpoint, the shape in current model is torch.Size([2050, 2560]).
        size mismatch for model.decoder.layers.0.self_attn.k_proj.weight: copying a param with shape torch.Size([0]) from checkpoint, the shape in current model is torch.Size([2560, 2560]).
        size mismatch for model.decoder.layers.0.self_attn.v_proj.weight: copying a param with shape torch.Size([0]) from checkpoint, the shape in current model is torch.Size([2560, 2560]).
        size mismatch for model.decoder.layers.0.self_attn.q_proj.weight: copying a param with shape torch.Size([0]) from checkpoint, the shape in current model is torch.Size([2560, 2560]).
        size mismatch for model.decoder.layers.0.self_attn.out_proj.weight: copying a param with shape torch.Size([0]) from checkpoint, the shape in current model is torch.Size([2560, 2560]).
        size mismatch for model.decoder.layers.0.fc1.weight: copying a param with shape torch.Size([0]) from checkpoint, the shape in current model is torch.Size([10240, 2560]).
        size mismatch for model.decoder.layers.0.fc2.weight: copying a param with shape torch.Size([0]) from checkpoint, the shape in current model is torch.Size([2560, 10240]).
        size mismatch for model.decoder.layers.1.self_attn.k_proj.weight: copying a param with shape torch.Size([0]) from checkpoint, the shape in current model is torch.Size([2560, 2560]).
        size mismatch for model.decoder.layers.1.self_attn.v_proj.weight: copying a param with shape torch.Size([0]) from checkpoint, the shape in current model is torch.Size([2560, 2560]).
        size mismatch for model.decoder.layers.1.self_attn.q_proj.weight: copying a param with shape torch.Size([0]) from checkpoint, the shape in current model is torch.Size([2560, 2560]).
        size mismatch for model.decoder.layers.1.self_attn.out_proj.weight: copying a param with shape torch.Size([0]) from checkpoint, the shape in current model is torch.Size([2560, 2560]).
        size mismatch for model.decoder.layers.1.fc1.weight: copying a param with shape torch.Size([0]) from checkpoint, the shape in current model is torch.Size([10240, 2560]).
        size mismatch for model.decoder.layers.1.fc2.weight: copying a param with shape torch.Size([0]) from checkpoint, the shape in current model is torch.Size([2560, 10240]).
        size mismatch for model.decoder.layers.2.self_attn.k_proj.weight: copying a param with shape torch.Size([0]) from checkpoint, the shape in current model is torch.Size([2560, 2560]).
        size mismatch for model.decoder.layers.2.self_attn.v_proj.weight: copying a param with shape torch.Size([0]) from checkpoint, the shape in current model is torch.Size([2560, 2560]).
        size mismatch for model.decoder.layers.2.self_attn.q_proj.weight: copying a param with shape torch.Size([0]) from checkpoint, the shape in current model is torch.Size([2560, 2560]).
        size mismatch for model.decoder.layers.2.self_attn.out_proj.weight: copying a param with shape torch.Size([0]) from checkpoint, the shape in current model is torch.Size([2560, 2560]).
        size mismatch for model.decoder.layers.2.fc1.weight: copying a param with shape torch.Size([0]) from checkpoint, the shape in current model is torch.Size([10240, 2560]).
        size mismatch for model.decoder.layers.2.fc2.weight: copying a param with shape torch.Size([0]) from checkpoint, the shape in current model is torch.Size([2560, 10240]).
        size mismatch for model.decoder.layers.3.self_attn.k_proj.weight: copying a param with shape torch.Size([0]) from checkpoint, the shape in current model is torch.Size([2560, 2560]).
        size mismatch for model.decoder.layers.3.self_attn.v_proj.weight: copying a param with shape torch.Size([0]) from checkpoint, the shape in current model is torch.Size([2560, 2560]).
        size mismatch for model.decoder.layers.3.self_attn.q_proj.weight: copying a param with shape torch.Size([0]) from checkpoint, the shape in current model is torch.Size([2560, 2560]).
        size mismatch for model.decoder.layers.3.self_attn.out_proj.weight: copying a param with shape torch.Size([0]) from checkpoint, the shape in current model is torch.Size([2560, 2560]).
        size mismatch for model.decoder.layers.3.fc1.weight: copying a param with shape torch.Size([0]) from checkpoint, the shape in current model is torch.Size([10240, 2560]).
        size mismatch for model.decoder.layers.3.fc2.weight: copying a param with shape torch.Size([0]) from checkpoint, the shape in current model is torch.Size([2560, 10240]).
        size mismatch for model.decoder.layers.4.self_attn.k_proj.weight: copying a param with shape torch.Size([0]) from checkpoint, the shape in current model is torch.Size([2560, 2560]).
        size mismatch for model.decoder.layers.4.self_attn.v_proj.weight: copying a param with shape torch.Size([0]) from checkpoint, the shape in current model is torch.Size([2560, 2560]).
        size mismatch for model.decoder.layers.4.self_attn.q_proj.weight: copying a param with shape torch.Size([0]) from checkpoint, the shape in current model is torch.Size([2560, 2560]).
        size mismatch for model.decoder.layers.4.self_attn.out_proj.weight: copying a param with shape torch.Size([0]) from checkpoint, the shape in current model is torch.Size([2560, 2560]).
        size mismatch for model.decoder.layers.4.fc1.weight: copying a param with shape torch.Size([0]) from checkpoint, the shape in current model is torch.Size([10240, 2560]).
        size mismatch for model.decoder.layers.4.fc2.weight: copying a param with shape torch.Size([0]) from checkpoint, the shape in current model is torch.Size([2560, 10240]).
        size mismatch for model.decoder.layers.5.self_attn.k_proj.weight: copying a param with shape torch.Size([0]) from checkpoint, the shape in current model is torch.Size([2560, 2560]).
        size mismatch for model.decoder.layers.5.self_attn.v_proj.weight: copying a param with shape torch.Size([0]) from checkpoint, the shape in current model is torch.Size([2560, 2560]).
        size mismatch for model.decoder.layers.5.self_attn.q_proj.weight: copying a param with shape torch.Size([0]) from checkpoint, the shape in current model is torch.Size([2560, 2560]).
        size mismatch for model.decoder.layers.5.self_attn.out_proj.weight: copying a param with shape torch.Size([0]) from checkpoint, the shape in current model is torch.Size([2560, 2560]).
        size mismatch for model.decoder.layers.5.fc1.weight: copying a param with shape torch.Size([0]) from checkpoint, the shape in current model is torch.Size([10240, 2560]).
        size mismatch for model.decoder.layers.5.fc2.weight: copying a param with shape torch.Size([0]) from checkpoint, the shape in current model is torch.Size([2560, 10240]).
        size mismatch for model.decoder.layers.6.self_attn.k_proj.weight: copying a param with shape torch.Size([0]) from checkpoint, the shape in current model is torch.Size([2560, 2560]).
        size mismatch for model.decoder.layers.6.self_attn.v_proj.weight: copying a param with shape torch.Size([0]) from checkpoint, the shape in current model is torch.Size([2560, 2560]).
        size mismatch for model.decoder.layers.6.self_attn.q_proj.weight: copying a param with shape torch.Size([0]) from checkpoint, the shape in current model is torch.Size([2560, 2560]).
        size mismatch for model.decoder.layers.6.self_attn.out_proj.weight: copying a param with shape torch.Size([0]) from checkpoint, the shape in current model is torch.Size([2560, 2560]).
        size mismatch for model.decoder.layers.6.fc1.weight: copying a param with shape torch.Size([0]) from checkpoint, the shape in current model is torch.Size([10240, 2560]).
        size mismatch for model.decoder.layers.6.fc2.weight: copying a param with shape torch.Size([0]) from checkpoint, the shape in current model is torch.Size([2560, 10240]).
        size mismatch for model.decoder.layers.7.self_attn.k_proj.weight: copying a param with shape torch.Size([0]) from checkpoint, the shape in current model is torch.Size([2560, 2560]).
        size mismatch for model.decoder.layers.7.self_attn.v_proj.weight: copying a param with shape torch.Size([0]) from checkpoint, the shape in current model is torch.Size([2560, 2560]).
        size mismatch for model.decoder.layers.7.self_attn.q_proj.weight: copying a param with shape torch.Size([0]) from checkpoint, the shape in current model is torch.Size([2560, 2560]).
        size mismatch for model.decoder.layers.7.self_attn.out_proj.weight: copying a param with shape torch.Size([0]) from checkpoint, the shape in current model is torch.Size([2560, 2560]).
        size mismatch for model.decoder.layers.7.fc1.weight: copying a param with shape torch.Size([0]) from checkpoint, the shape in current model is torch.Size([10240, 2560]).
        size mismatch for model.decoder.layers.7.fc2.weight: copying a param with shape torch.Size([0]) from checkpoint, the shape in current model is torch.Size([2560, 10240]).
        size mismatch for model.decoder.layers.8.self_attn.k_proj.weight: copying a param with shape torch.Size([0]) from checkpoint, the shape in current model is torch.Size([2560, 2560]).
        size mismatch for model.decoder.layers.8.self_attn.v_proj.weight: copying a param with shape torch.Size([0]) from checkpoint, the shape in current model is torch.Size([2560, 2560]).
        size mismatch for model.decoder.layers.8.self_attn.q_proj.weight: copying a param with shape torch.Size([0]) from checkpoint, the shape in current model is torch.Size([2560, 2560]).
        size mismatch for model.decoder.layers.8.self_attn.out_proj.weight: copying a param with shape torch.Size([0]) from checkpoint, the shape in current model is torch.Size([2560, 2560]).
        size mismatch for model.decoder.layers.8.fc1.weight: copying a param with shape torch.Size([0]) from checkpoint, the shape in current model is torch.Size([10240, 2560]).
        size mismatch for model.decoder.layers.8.fc2.weight: copying a param with shape torch.Size([0]) from checkpoint, the shape in current model is torch.Size([2560, 10240]).
        size mismatch for model.decoder.layers.9.self_attn.k_proj.weight: copying a param with shape torch.Size([0]) from checkpoint, the shape in current model is torch.Size([2560, 2560]).
        size mismatch for model.decoder.layers.9.self_attn.v_proj.weight: copying a param with shape torch.Size([0]) from checkpoint, the shape in current model is torch.Size([2560, 2560]).
        size mismatch for model.decoder.layers.9.self_attn.q_proj.weight: copying a param with shape torch.Size([0]) from checkpoint, the shape in current model is torch.Size([2560, 2560]).
        size mismatch for model.decoder.layers.9.self_attn.out_proj.weight: copying a param with shape torch.Size([0]) from checkpoint, the shape in current model is torch.Size([2560, 2560]).
        size mismatch for model.decoder.layers.9.fc1.weight: copying a param with shape torch.Size([0]) from checkpoint, the shape in current model is torch.Size([10240, 2560]).
        size mismatch for model.decoder.layers.9.fc2.weight: copying a param with shape torch.Size([0]) from checkpoint, the shape in current model is torch.Size([2560, 10240]).
        size mismatch for model.decoder.layers.10.self_attn.k_proj.weight: copying a param with shape torch.Size([0]) from checkpoint, the shape in current model is torch.Size([2560, 2560]).
        size mismatch for model.decoder.layers.10.self_attn.v_proj.weight: copying a param with shape torch.Size([0]) from checkpoint, the shape in current model is torch.Size([2560, 2560]).
        size mismatch for model.decoder.layers.10.self_attn.q_proj.weight: copying a param with shape torch.Size([0]) from checkpoint, the shape in current model is torch.Size([2560, 2560]).
        size mismatch for model.decoder.layers.10.self_attn.out_proj.weight: copying a param with shape torch.Size([0]) from checkpoint, the shape in current model is torch.Size([2560, 2560]).
        size mismatch for model.decoder.layers.10.fc1.weight: copying a param with shape torch.Size([0]) from checkpoint, the shape in current model is torch.Size([10240, 2560]).
        size mismatch for model.decoder.layers.10.fc2.weight: copying a param with shape torch.Size([0]) from checkpoint, the shape in current model is torch.Size([2560, 10240]).
        size mismatch for model.decoder.layers.11.self_attn.k_proj.weight: copying a param with shape torch.Size([0]) from checkpoint, the shape in current model is torch.Size([2560, 2560]).
        size mismatch for model.decoder.layers.11.self_attn.v_proj.weight: copying a param with shape torch.Size([0]) from checkpoint, the shape in current model is torch.Size([2560, 2560]).
        size mismatch for model.decoder.layers.11.self_attn.q_proj.weight: copying a param with shape torch.Size([0]) from checkpoint, the shape in current model is torch.Size([2560, 2560]).
        size mismatch for model.decoder.layers.11.self_attn.out_proj.weight: copying a param with shape torch.Size([0]) from checkpoint, the shape in current model is torch.Size([2560, 2560]).
        size mismatch for model.decoder.layers.11.fc1.weight: copying a param with shape torch.Size([0]) from checkpoint, the shape in current model is torch.Size([10240, 2560]).
        size mismatch for model.decoder.layers.11.fc2.weight: copying a param with shape torch.Size([0]) from checkpoint, the shape in current model is torch.Size([2560, 10240]).
        size mismatch for model.decoder.layers.12.self_attn.k_proj.weight: copying a param with shape torch.Size([0]) from checkpoint, the shape in current model is torch.Size([2560, 2560]).
        size mismatch for model.decoder.layers.12.self_attn.v_proj.weight: copying a param with shape torch.Size([0]) from checkpoint, the shape in current model is torch.Size([2560, 2560]).
        size mismatch for model.decoder.layers.12.self_attn.q_proj.weight: copying a param with shape torch.Size([0]) from checkpoint, the shape in current model is torch.Size([2560, 2560]).
        size mismatch for model.decoder.layers.12.self_attn.out_proj.weight: copying a param with shape torch.Size([0]) from checkpoint, the shape in current model is torch.Size([2560, 2560]).
        size mismatch for model.decoder.layers.12.fc1.weight: copying a param with shape torch.Size([0]) from checkpoint, the shape in current model is torch.Size([10240, 2560]).
        size mismatch for model.decoder.layers.12.fc2.weight: copying a param with shape torch.Size([0]) from checkpoint, the shape in current model is torch.Size([2560, 10240]).
        size mismatch for model.decoder.layers.13.self_attn.k_proj.weight: copying a param with shape torch.Size([0]) from checkpoint, the shape in current model is torch.Size([2560, 2560]).
        size mismatch for model.decoder.layers.13.self_attn.v_proj.weight: copying a param with shape torch.Size([0]) from checkpoint, the shape in current model is torch.Size([2560, 2560]).
        size mismatch for model.decoder.layers.13.self_attn.q_proj.weight: copying a param with shape torch.Size([0]) from checkpoint, the shape in current model is torch.Size([2560, 2560]).
        size mismatch for model.decoder.layers.13.self_attn.out_proj.weight: copying a param with shape torch.Size([0]) from checkpoint, the shape in current model is torch.Size([2560, 2560]).
        size mismatch for model.decoder.layers.13.fc1.weight: copying a param with shape torch.Size([0]) from checkpoint, the shape in current model is torch.Size([10240, 2560]).
        size mismatch for model.decoder.layers.13.fc2.weight: copying a param with shape torch.Size([0]) from checkpoint, the shape in current model is torch.Size([2560, 10240]).
        size mismatch for model.decoder.layers.14.self_attn.k_proj.weight: copying a param with shape torch.Size([0]) from checkpoint, the shape in current model is torch.Size([2560, 2560]).
        size mismatch for model.decoder.layers.14.self_attn.v_proj.weight: copying a param with shape torch.Size([0]) from checkpoint, the shape in current model is torch.Size([2560, 2560]).
        size mismatch for model.decoder.layers.14.self_attn.q_proj.weight: copying a param with shape torch.Size([0]) from checkpoint, the shape in current model is torch.Size([2560, 2560]).
        size mismatch for model.decoder.layers.14.self_attn.out_proj.weight: copying a param with shape torch.Size([0]) from checkpoint, the shape in current model is torch.Size([2560, 2560]).
        size mismatch for model.decoder.layers.14.fc1.weight: copying a param with shape torch.Size([0]) from checkpoint, the shape in current model is torch.Size([10240, 2560]).
        size mismatch for model.decoder.layers.14.fc2.weight: copying a param with shape torch.Size([0]) from checkpoint, the shape in current model is torch.Size([2560, 10240]).
        size mismatch for model.decoder.layers.15.self_attn.k_proj.weight: copying a param with shape torch.Size([0]) from checkpoint, the shape in current model is torch.Size([2560, 2560]).
        size mismatch for model.decoder.layers.15.self_attn.v_proj.weight: copying a param with shape torch.Size([0]) from checkpoint, the shape in current model is torch.Size([2560, 2560]).
        size mismatch for model.decoder.layers.15.self_attn.q_proj.weight: copying a param with shape torch.Size([0]) from checkpoint, the shape in current model is torch.Size([2560, 2560]).
        size mismatch for model.decoder.layers.15.self_attn.out_proj.weight: copying a param with shape torch.Size([0]) from checkpoint, the shape in current model is torch.Size([2560, 2560]).
        size mismatch for model.decoder.layers.15.fc1.weight: copying a param with shape torch.Size([0]) from checkpoint, the shape in current model is torch.Size([10240, 2560]).
        size mismatch for model.decoder.layers.15.fc2.weight: copying a param with shape torch.Size([0]) from checkpoint, the shape in current model is torch.Size([2560, 10240]).
        size mismatch for model.decoder.layers.16.self_attn.k_proj.weight: copying a param with shape torch.Size([0]) from checkpoint, the shape in current model is torch.Size([2560, 2560]).
        size mismatch for model.decoder.layers.16.self_attn.v_proj.weight: copying a param with shape torch.Size([0]) from checkpoint, the shape in current model is torch.Size([2560, 2560]).
        size mismatch for model.decoder.layers.16.self_attn.q_proj.weight: copying a param with shape torch.Size([0]) from checkpoint, the shape in current model is torch.Size([2560, 2560]).
        size mismatch for model.decoder.layers.16.self_attn.out_proj.weight: copying a param with shape torch.Size([0]) from checkpoint, the shape in current model is torch.Size([2560, 2560]).
        size mismatch for model.decoder.layers.16.fc1.weight: copying a param with shape torch.Size([0]) from checkpoint, the shape in current model is torch.Size([10240, 2560]).
        size mismatch for model.decoder.layers.16.fc2.weight: copying a param with shape torch.Size([0]) from checkpoint, the shape in current model is torch.Size([2560, 10240]).
        size mismatch for model.decoder.layers.17.self_attn.k_proj.weight: copying a param with shape torch.Size([0]) from checkpoint, the shape in current model is torch.Size([2560, 2560]).
        size mismatch for model.decoder.layers.17.self_attn.v_proj.weight: copying a param with shape torch.Size([0]) from checkpoint, the shape in current model is torch.Size([2560, 2560]).
        size mismatch for model.decoder.layers.17.self_attn.q_proj.weight: copying a param with shape torch.Size([0]) from checkpoint, the shape in current model is torch.Size([2560, 2560]).
        size mismatch for model.decoder.layers.17.self_attn.out_proj.weight: copying a param with shape torch.Size([0]) from checkpoint, the shape in current model is torch.Size([2560, 2560]).
        size mismatch for model.decoder.layers.17.fc1.weight: copying a param with shape torch.Size([0]) from checkpoint, the shape in current model is torch.Size([10240, 2560]).
        size mismatch for model.decoder.layers.17.fc2.weight: copying a param with shape torch.Size([0]) from checkpoint, the shape in current model is torch.Size([2560, 10240]).
        size mismatch for model.decoder.layers.18.self_attn.k_proj.weight: copying a param with shape torch.Size([0]) from checkpoint, the shape in current model is torch.Size([2560, 2560]).
        size mismatch for model.decoder.layers.18.self_attn.v_proj.weight: copying a param with shape torch.Size([0]) from checkpoint, the shape in current model is torch.Size([2560, 2560]).
        size mismatch for model.decoder.layers.18.self_attn.q_proj.weight: copying a param with shape torch.Size([0]) from checkpoint, the shape in current model is torch.Size([2560, 2560]).
        size mismatch for model.decoder.layers.18.self_attn.out_proj.weight: copying a param with shape torch.Size([0]) from checkpoint, the shape in current model is torch.Size([2560, 2560]).
        size mismatch for model.decoder.layers.18.fc1.weight: copying a param with shape torch.Size([0]) from checkpoint, the shape in current model is torch.Size([10240, 2560]).
        size mismatch for model.decoder.layers.18.fc2.weight: copying a param with shape torch.Size([0]) from checkpoint, the shape in current model is torch.Size([2560, 10240]).
        size mismatch for model.decoder.layers.19.self_attn.k_proj.weight: copying a param with shape torch.Size([0]) from checkpoint, the shape in current model is torch.Size([2560, 2560]).
        size mismatch for model.decoder.layers.19.self_attn.v_proj.weight: copying a param with shape torch.Size([0]) from checkpoint, the shape in current model is torch.Size([2560, 2560]).
        size mismatch for model.decoder.layers.19.self_attn.q_proj.weight: copying a param with shape torch.Size([0]) from checkpoint, the shape in current model is torch.Size([2560, 2560]).
        size mismatch for model.decoder.layers.19.self_attn.out_proj.weight: copying a param with shape torch.Size([0]) from checkpoint, the shape in current model is torch.Size([2560, 2560]).
        size mismatch for model.decoder.layers.19.fc1.weight: copying a param with shape torch.Size([0]) from checkpoint, the shape in current model is torch.Size([10240, 2560]).
        size mismatch for model.decoder.layers.19.fc2.weight: copying a param with shape torch.Size([0]) from checkpoint, the shape in current model is torch.Size([2560, 10240]).
        size mismatch for model.decoder.layers.20.self_attn.k_proj.weight: copying a param with shape torch.Size([0]) from checkpoint, the shape in current model is torch.Size([2560, 2560]).
        size mismatch for model.decoder.layers.20.self_attn.v_proj.weight: copying a param with shape torch.Size([0]) from checkpoint, the shape in current model is torch.Size([2560, 2560]).
        size mismatch for model.decoder.layers.20.self_attn.q_proj.weight: copying a param with shape torch.Size([0]) from checkpoint, the shape in current model is torch.Size([2560, 2560]).
        size mismatch for model.decoder.layers.20.self_attn.out_proj.weight: copying a param with shape torch.Size([0]) from checkpoint, the shape in current model is torch.Size([2560, 2560]).
        size mismatch for model.decoder.layers.20.fc1.weight: copying a param with shape torch.Size([0]) from checkpoint, the shape in current model is torch.Size([10240, 2560]).
        size mismatch for model.decoder.layers.20.fc2.weight: copying a param with shape torch.Size([0]) from checkpoint, the shape in current model is torch.Size([2560, 10240]).
        size mismatch for model.decoder.layers.21.self_attn.k_proj.weight: copying a param with shape torch.Size([0]) from checkpoint, the shape in current model is torch.Size([2560, 2560]).
        size mismatch for model.decoder.layers.21.self_attn.v_proj.weight: copying a param with shape torch.Size([0]) from checkpoint, the shape in current model is torch.Size([2560, 2560]).
        size mismatch for model.decoder.layers.21.self_attn.q_proj.weight: copying a param with shape torch.Size([0]) from checkpoint, the shape in current model is torch.Size([2560, 2560]).
        size mismatch for model.decoder.layers.21.self_attn.out_proj.weight: copying a param with shape torch.Size([0]) from checkpoint, the shape in current model is torch.Size([2560, 2560]).
        size mismatch for model.decoder.layers.21.fc1.weight: copying a param with shape torch.Size([0]) from checkpoint, the shape in current model is torch.Size([10240, 2560]).
        size mismatch for model.decoder.layers.21.fc2.weight: copying a param with shape torch.Size([0]) from checkpoint, the shape in current model is torch.Size([2560, 10240]).
        size mismatch for model.decoder.layers.22.self_attn.k_proj.weight: copying a param with shape torch.Size([0]) from checkpoint, the shape in current model is torch.Size([2560, 2560]).
        size mismatch for model.decoder.layers.22.self_attn.v_proj.weight: copying a param with shape torch.Size([0]) from checkpoint, the shape in current model is torch.Size([2560, 2560]).
        size mismatch for model.decoder.layers.22.self_attn.q_proj.weight: copying a param with shape torch.Size([0]) from checkpoint, the shape in current model is torch.Size([2560, 2560]).
        size mismatch for model.decoder.layers.22.self_attn.out_proj.weight: copying a param with shape torch.Size([0]) from checkpoint, the shape in current model is torch.Size([2560, 2560]).
        size mismatch for model.decoder.layers.22.fc1.weight: copying a param with shape torch.Size([0]) from checkpoint, the shape in current model is torch.Size([10240, 2560]).
        size mismatch for model.decoder.layers.22.fc2.weight: copying a param with shape torch.Size([0]) from checkpoint, the shape in current model is torch.Size([2560, 10240]).
        size mismatch for model.decoder.layers.23.self_attn.k_proj.weight: copying a param with shape torch.Size([0]) from checkpoint, the shape in current model is torch.Size([2560, 2560]).
        size mismatch for model.decoder.layers.23.self_attn.v_proj.weight: copying a param with shape torch.Size([0]) from checkpoint, the shape in current model is torch.Size([2560, 2560]).
        size mismatch for model.decoder.layers.23.self_attn.q_proj.weight: copying a param with shape torch.Size([0]) from checkpoint, the shape in current model is torch.Size([2560, 2560]).
        size mismatch for model.decoder.layers.23.self_attn.out_proj.weight: copying a param with shape torch.Size([0]) from checkpoint, the shape in current model is torch.Size([2560, 2560]).
        size mismatch for model.decoder.layers.23.fc1.weight: copying a param with shape torch.Size([0]) from checkpoint, the shape in current model is torch.Size([10240, 2560]).
        size mismatch for model.decoder.layers.23.fc2.weight: copying a param with shape torch.Size([0]) from checkpoint, the shape in current model is torch.Size([2560, 10240]).
        size mismatch for model.decoder.layers.24.self_attn.k_proj.weight: copying a param with shape torch.Size([0]) from checkpoint, the shape in current model is torch.Size([2560, 2560]).
        size mismatch for model.decoder.layers.24.self_attn.v_proj.weight: copying a param with shape torch.Size([0]) from checkpoint, the shape in current model is torch.Size([2560, 2560]).
        size mismatch for model.decoder.layers.24.self_attn.q_proj.weight: copying a param with shape torch.Size([0]) from checkpoint, the shape in current model is torch.Size([2560, 2560]).
        size mismatch for model.decoder.layers.24.self_attn.out_proj.weight: copying a param with shape torch.Size([0]) from checkpoint, the shape in current model is torch.Size([2560, 2560]).
        size mismatch for model.decoder.layers.24.fc1.weight: copying a param with shape torch.Size([0]) from checkpoint, the shape in current model is torch.Size([10240, 2560]).
        size mismatch for model.decoder.layers.24.fc2.weight: copying a param with shape torch.Size([0]) from checkpoint, the shape in current model is torch.Size([2560, 10240]).
        size mismatch for model.decoder.layers.25.self_attn.k_proj.weight: copying a param with shape torch.Size([0]) from checkpoint, the shape in current model is torch.Size([2560, 2560]).
        size mismatch for model.decoder.layers.25.self_attn.v_proj.weight: copying a param with shape torch.Size([0]) from checkpoint, the shape in current model is torch.Size([2560, 2560]).
        size mismatch for model.decoder.layers.25.self_attn.q_proj.weight: copying a param with shape torch.Size([0]) from checkpoint, the shape in current model is torch.Size([2560, 2560]).
        size mismatch for model.decoder.layers.25.self_attn.out_proj.weight: copying a param with shape torch.Size([0]) from checkpoint, the shape in current model is torch.Size([2560, 2560]).
        size mismatch for model.decoder.layers.25.fc1.weight: copying a param with shape torch.Size([0]) from checkpoint, the shape in current model is torch.Size([10240, 2560]).
        size mismatch for model.decoder.layers.25.fc2.weight: copying a param with shape torch.Size([0]) from checkpoint, the shape in current model is torch.Size([2560, 10240]).
        size mismatch for model.decoder.layers.26.self_attn.k_proj.weight: copying a param with shape torch.Size([0]) from checkpoint, the shape in current model is torch.Size([2560, 2560]).
        size mismatch for model.decoder.layers.26.self_attn.v_proj.weight: copying a param with shape torch.Size([0]) from checkpoint, the shape in current model is torch.Size([2560, 2560]).
        size mismatch for model.decoder.layers.26.self_attn.q_proj.weight: copying a param with shape torch.Size([0]) from checkpoint, the shape in current model is torch.Size([2560, 2560]).
        size mismatch for model.decoder.layers.26.self_attn.out_proj.weight: copying a param with shape torch.Size([0]) from checkpoint, the shape in current model is torch.Size([2560, 2560]).
        size mismatch for model.decoder.layers.26.fc1.weight: copying a param with shape torch.Size([0]) from checkpoint, the shape in current model is torch.Size([10240, 2560]).
        size mismatch for model.decoder.layers.26.fc2.weight: copying a param with shape torch.Size([0]) from checkpoint, the shape in current model is torch.Size([2560, 10240]).
        size mismatch for model.decoder.layers.27.self_attn.k_proj.weight: copying a param with shape torch.Size([0]) from checkpoint, the shape in current model is torch.Size([2560, 2560]).
        size mismatch for model.decoder.layers.27.self_attn.v_proj.weight: copying a param with shape torch.Size([0]) from checkpoint, the shape in current model is torch.Size([2560, 2560]).
        size mismatch for model.decoder.layers.27.self_attn.q_proj.weight: copying a param with shape torch.Size([0]) from checkpoint, the shape in current model is torch.Size([2560, 2560]).
        size mismatch for model.decoder.layers.27.self_attn.out_proj.weight: copying a param with shape torch.Size([0]) from checkpoint, the shape in current model is torch.Size([2560, 2560]).
        size mismatch for model.decoder.layers.27.fc1.weight: copying a param with shape torch.Size([0]) from checkpoint, the shape in current model is torch.Size([10240, 2560]).
        size mismatch for model.decoder.layers.27.fc2.weight: copying a param with shape torch.Size([0]) from checkpoint, the shape in current model is torch.Size([2560, 10240]).
        size mismatch for model.decoder.layers.28.self_attn.k_proj.weight: copying a param with shape torch.Size([0]) from checkpoint, the shape in current model is torch.Size([2560, 2560]).
        size mismatch for model.decoder.layers.28.self_attn.v_proj.weight: copying a param with shape torch.Size([0]) from checkpoint, the shape in current model is torch.Size([2560, 2560]).
        size mismatch for model.decoder.layers.28.self_attn.q_proj.weight: copying a param with shape torch.Size([0]) from checkpoint, the shape in current model is torch.Size([2560, 2560]).
        size mismatch for model.decoder.layers.28.self_attn.out_proj.weight: copying a param with shape torch.Size([0]) from checkpoint, the shape in current model is torch.Size([2560, 2560]).
        size mismatch for model.decoder.layers.28.fc1.weight: copying a param with shape torch.Size([0]) from checkpoint, the shape in current model is torch.Size([10240, 2560]).
        size mismatch for model.decoder.layers.28.fc2.weight: copying a param with shape torch.Size([0]) from checkpoint, the shape in current model is torch.Size([2560, 10240]).
        size mismatch for model.decoder.layers.29.self_attn.k_proj.weight: copying a param with shape torch.Size([0]) from checkpoint, the shape in current model is torch.Size([2560, 2560]).
        size mismatch for model.decoder.layers.29.self_attn.v_proj.weight: copying a param with shape torch.Size([0]) from checkpoint, the shape in current model is torch.Size([2560, 2560]).
        size mismatch for model.decoder.layers.29.self_attn.q_proj.weight: copying a param with shape torch.Size([0]) from checkpoint, the shape in current model is torch.Size([2560, 2560]).
        size mismatch for model.decoder.layers.29.self_attn.out_proj.weight: copying a param with shape torch.Size([0]) from checkpoint, the shape in current model is torch.Size([2560, 2560]).
        size mismatch for model.decoder.layers.29.fc1.weight: copying a param with shape torch.Size([0]) from checkpoint, the shape in current model is torch.Size([10240, 2560]).
        size mismatch for model.decoder.layers.29.fc2.weight: copying a param with shape torch.Size([0]) from checkpoint, the shape in current model is torch.Size([2560, 10240]).
        size mismatch for model.decoder.layers.30.self_attn.k_proj.weight: copying a param with shape torch.Size([0]) from checkpoint, the shape in current model is torch.Size([2560, 2560]).
        size mismatch for model.decoder.layers.30.self_attn.v_proj.weight: copying a param with shape torch.Size([0]) from checkpoint, the shape in current model is torch.Size([2560, 2560]).
        size mismatch for model.decoder.layers.30.self_attn.q_proj.weight: copying a param with shape torch.Size([0]) from checkpoint, the shape in current model is torch.Size([2560, 2560]).
        size mismatch for model.decoder.layers.30.self_attn.out_proj.weight: copying a param with shape torch.Size([0]) from checkpoint, the shape in current model is torch.Size([2560, 2560]).
        size mismatch for model.decoder.layers.30.fc1.weight: copying a param with shape torch.Size([0]) from checkpoint, the shape in current model is torch.Size([10240, 2560]).
        size mismatch for model.decoder.layers.30.fc2.weight: copying a param with shape torch.Size([0]) from checkpoint, the shape in current model is torch.Size([2560, 10240]).
        size mismatch for model.decoder.layers.31.self_attn.k_proj.weight: copying a param with shape torch.Size([0]) from checkpoint, the shape in current model is torch.Size([2560, 2560]).
        size mismatch for model.decoder.layers.31.self_attn.v_proj.weight: copying a param with shape torch.Size([0]) from checkpoint, the shape in current model is torch.Size([2560, 2560]).
        size mismatch for model.decoder.layers.31.self_attn.q_proj.weight: copying a param with shape torch.Size([0]) from checkpoint, the shape in current model is torch.Size([2560, 2560]).
        size mismatch for model.decoder.layers.31.self_attn.out_proj.weight: copying a param with shape torch.Size([0]) from checkpoint, the shape in current model is torch.Size([2560, 2560]).
        size mismatch for model.decoder.layers.31.fc1.weight: copying a param with shape torch.Size([0]) from checkpoint, the shape in current model is torch.Size([10240, 2560]).
        size mismatch for model.decoder.layers.31.fc2.weight: copying a param with shape torch.Size([0]) from checkpoint, the shape in current model is torch.Size([2560, 10240]).
        size mismatch for lm_head.weight: copying a param with shape torch.Size([0]) from checkpoint, the shape in current model is torch.Size([50272, 2560]).
        You may consider adding `ignore_mismatched_sizes=True` in the model `from_pretrained` method.
```<|||||>cc @pacman100 <|||||>Hello @JohnnyRacer, can you provide script of how you are saving and loading weights post training? Also, what is the accelerate config? <|||||>Hey @pacman100 .  I am using the unmodified `run_clm_no_trainer.py` from the latest commit of `transformers`, with the following training commands : 
```
accelerate launch run_clm_no_trainer.py \                                                                
--model_name_or_path facebook/opt-1.3b \
--dataset_name  wikitext \
--num_train_epochs 6 \
--block_size 128 \
--output_dir ./opt-1.3b-wikitext
```
Accelerate config from config file : 
```json
{
  "compute_environment": "LOCAL_MACHINE",
  "deepspeed_config": {
    "gradient_accumulation_steps": 1,
    "offload_optimizer_device": "cpu",
    "offload_param_device": "none",
    "zero3_init_flag": false,
    "zero3_save_16bit_model": false,
    "zero_stage": 3
  },
  "distributed_type": "DEEPSPEED",
  "downcast_bf16": "no",
  "fsdp_config": {},
  "gpu_ids": null,
  "machine_rank": 0,
  "main_process_ip": null,
  "main_process_port": null,
  "main_training_function": "main",
  "mixed_precision": "fp16",
  "num_machines": 1,
  "num_processes": 2,
  "rdzv_backend": "static",
  "same_network": true,
  "use_cpu": false
}
```
The config doesn't show it but I have it configured for a multi-GPU setup on a single local instance with 2 accelerators.

And I have tried loading the models with the following commands using `transfomers.AutoModelForCausalLM`

```python
import torch
from transformers import AutoModelForCausalLM
model = AutoModelForCausalLM.from_pretrained("facebook/opt-1.3b")
model.load_state_dict(torch.load("./opt-1.3b-wikitext/pytorch_model.bin"))
```
As well as loading the model directly from the directory via the model config by using:
```python
from transformers import AutoModelForCausalLM
model = AutoModelForCausalLM.from_pretrained("opt-1.3b-wikitext")
```
Both of these methods results in the same error that I have described above.<|||||>Hello @JohnnyRacer , please refer below code snippet on changes required when saving deepspeed ZeRO-3 model. The example can be found here: [deepspeed_with_config_support.py](https://github.com/huggingface/accelerate/blob/main/examples/by_feature/deepspeed_with_config_support.py)

https://github.com/huggingface/accelerate/blob/cea6aaa1161d45f7f23ef33fcc3b0a5999ebb5a1/examples/by_feature/deepspeed_with_config_support.py#L712-L723<|||||>Thanks @pacman100 . Just finished training the model and can confirm loading works correctly with the script you have linked. However I still had to modify the script to include this fix for an issue I had earlier to ensure the weights can be correctly loaded, [link to issue](https://github.com/huggingface/transformers/issues/19959). <|||||>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.