Model saving have dimension issue while using deepspeed stage 3 with multi node for larger models
### System Info

Python 3.8.10
transformers 4.30.2
accelerate 0.20.3
deepspeed 0.9.5

### Who can help?

@pacman100 

### Information

- [ ] The official example scripts
- [ ] My own modified scripts

### Tasks

- [ ] An officially supported task in the `examples` folder (such as GLUE/SQuAD, ...)
- [ ] My own task or dataset (give details below)

### Reproduction

I'm finetuning LLaMa30B using the command below.

`accelerate launch  src/train_sft.py --model_name_or_path huggyllama/llama-30b --do_train --dataset dummy_identity --finetuning_type full --output_dir output/30B-sft-identity-v1 --overwrite_cache --per_device_train_batch_size 4 --gradient_accumulation_steps 4 --lr_scheduler_type cosine --logging_steps 10 --save_steps 200 --learning_rate 5e-5 --num_train_epochs 3 --plot_loss --fp16 --deepspeed ds_config.json --report_to wandb`

I'm running this on 2 nodes with 4 A100 80GB on each. For deepspeed I'm using stage 3. When I tried to load the saved model it's giving the error as below. But this issue is not there for LLaMa 7B. So I assume it has something to do with the stage 3 optimization and gathering for saving.

<img width="1156" alt="image" src="https://github.com/huggingface/transformers/assets/12937285/d4f1b234-b67b-457a-84b3-2f64c6101418">

This is my deepspeed config

```
{  "fp16": {      "enabled": "auto",      "loss_scale": 0,
      "loss_scale_window": 1000,
      "initial_scale_power": 16,
      "hysteresis": 2,
      "min_loss_scale": 1
  },

  "optimizer": {
      "type": "AdamW",
      "params": {
          "lr": "auto",
          "betas": "auto",
          "eps": "auto",
          "weight_decay": "auto"
      }
  },

  "scheduler": {
      "type": "WarmupLR",
      "params": {
          "warmup_min_lr": "auto",
          "warmup_max_lr": "auto",
          "warmup_num_steps": "auto"
      }
  },

  "zero_optimization": {
      "stage": 3,
      "overlap_comm": true,
      "contiguous_gradients": true,
      "sub_group_size": 1e9,
      "reduce_bucket_size": "auto",
      "stage3_prefetch_bucket_size": "auto",
      "stage3_param_persistence_threshold": "auto",
      "stage3_max_live_parameters": 1e9,
      "stage3_max_reuse_distance": 1e9,
      "stage3_gather_16bit_weights_on_model_save": true
  },

  "gradient_accumulation_steps": "auto",
  "gradient_clipping": "auto",
  "steps_per_print": 2000,
  "train_batch_size": "auto",
  "train_micro_batch_size_per_gpu": "auto",
  "wall_clock_breakdown": false
}
```

A similar [issue](https://github.com/hiyouga/LLaMA-Efficient-Tuning/issues/70#issuecomment-1626876405) was reported on the open-source repo I'm using. I can see they were able to resolve it without using HF trainer and write a script with accelerate. 



### Expected behavior

The multinode stage 3 trained model should be saved without any errors. 
@dittops 
To be fair, it was said that it was resolved, but we have no means of validating that, since the actual code was not shared. 