Model saving have dimension issue while using deepspeed stage 3 with multi node for larger models
Author: dittops
Pull request: no

### 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. 

Comments:
@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. 

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.