[DeepSpeed] Convert from fp16 to fp32 issue zero_to_fp32.py
## Environment info
<!-- You can run the command `transformers-cli env` and copy-and-paste its output below.
     Don't forget to fill out the missing fields in that output! -->

- `transformers` version: 4.7.0
- Platform: Linux-4.4.0-1128-aws-x86_64-with-debian-stretch-sid
- Python version: 3.7.6
- PyTorch version (GPU?): 1.8.1+cu111 (True)
- Tensorflow version (GPU?): not installed (NA)
- Using GPU in script?: yes
- Using distributed or parallel set-up in script?: using deepspeed v 0.4.0

### Who can help
@stas00 thanks for opening new doors by showing how to train large transformer models. 

## Information

Model I am using (Bert, XLNet ...): t5-3b

The problem arises when using:
* [X] the official example scripts: (give details below)
* [ ] my own modified scripts: (give details below)

The tasks I am working on is:
* [ ] an official GLUE/SQUaD task: (give the name)
* [ ] my own task or dataset: (give details below)

## To reproduce

I trained a t5-3b model on A100-40GB single host using deepspeed using tips from https://github.com/huggingface/transformers/issues/9996 and huggingface docs. After a checkpoint was saved, I used the provided `zero_to_fp32.py` 

Steps to reproduce the behavior:

1. Using deepspeed **zero2 config** to train a **t5-3b** model.
2. Tried converting the deepspeed saved fp16 checkpoint (checkpoint-60000) to fp32
2. I went into the checkpoint-60000 dir and ran the provided command `python zero_to_fp32.py global_step60001 pytorch_model_fp32.bin` this is based on deepspeed documentation.
3. however I get the crash shown below. 

```python
python zero_to_fp32.py global_step60001 pytorch_model_fp32.bin

Processing zero checkpoint 'global_step60001'
Detected checkpoint of type zero stage 2, world_size: 8
Traceback (most recent call last):
  File "zero_to_fp32.py", line 170, in <module>
    convert_zero_chkpt_to_fp32_consolid_state_dict(args.checkpoint_dir, args.output_file)
  File "zero_to_fp32.py", line 128, in convert_zero_chkpt_to_fp32_consolid_state_dict
    unpartitioned_numel).view(shape)
RuntimeError: start (2499259392) + length (16777216) exceeds dimension size (2499738752).
```

contents inside `global_step60001` folder
```
total 34G
-rw-rw-r-- 1 ubuntu ubuntu 5.4G Jun 28 06:15 mp_rank_00_model_states.pt
-rw-rw-r-- 1 ubuntu ubuntu 4.0G Jun 28 06:15 zero_pp_rank_4_mp_rank_00_optim_states.pt
-rw-rw-r-- 1 ubuntu ubuntu 4.0G Jun 28 06:15 zero_pp_rank_2_mp_rank_00_optim_states.pt
-rw-rw-r-- 1 ubuntu ubuntu 4.0G Jun 28 06:15 zero_pp_rank_7_mp_rank_00_optim_states.pt
-rw-rw-r-- 1 ubuntu ubuntu 4.0G Jun 28 06:15 zero_pp_rank_3_mp_rank_00_optim_states.pt
-rw-rw-r-- 1 ubuntu ubuntu 4.0G Jun 28 06:15 zero_pp_rank_6_mp_rank_00_optim_states.pt
-rw-rw-r-- 1 ubuntu ubuntu 4.0G Jun 28 06:15 zero_pp_rank_0_mp_rank_00_optim_states.pt
-rw-rw-r-- 1 ubuntu ubuntu 4.0G Jun 28 06:15 zero_pp_rank_1_mp_rank_00_optim_states.pt
```
Oddly, I see one file not present "rank_2" I am assuming each GPU saves its optimizer state. But I have not modified any code to cause this issue. Please help!

## Expected behavior

<!-- A clear and concise description of what you would expect to happen. -->
Since I was running the provided code `zero_to_fp32.py` I would expect this to create the `fp32` model binary which does not happen due to the cryptic crash. 

Happy to provide more information.


Thank you for this excellent report, @srikar2097 

This is fascinating!

Where is `zero_pp_rank_5_mp_rank_00_optim_states.pt`? I presume this was an 8-gpu run, so we are missing one rank.

Could you re-run again and see if it was a fluke? Clearly one process failed to save its optimizer states.
<|||||>@stas00 When I looked at what got saved in the previous check-point (previous to what I shared in my original report), this is what I see:

```
ll ./checkpoint-50000/global_step50001/
total 34G
-rw-rw-r-- 1 ubuntu ubuntu 5.4G Jun 27 01:27 mp_rank_00_model_states.pt
-rw-rw-r-- 1 ubuntu ubuntu 4.0G Jun 27 01:27 zero_pp_rank_3_mp_rank_00_optim_states.pt
-rw-rw-r-- 1 ubuntu ubuntu 4.0G Jun 27 01:27 zero_pp_rank_1_mp_rank_00_optim_states.pt
-rw-rw-r-- 1 ubuntu ubuntu 4.0G Jun 27 01:27 zero_pp_rank_7_mp_rank_00_optim_states.pt
-rw-rw-r-- 1 ubuntu ubuntu 4.0G Jun 27 01:27 zero_pp_rank_4_mp_rank_00_optim_states.pt
-rw-rw-r-- 1 ubuntu ubuntu 4.0G Jun 27 01:27 zero_pp_rank_6_mp_rank_00_optim_states.pt
-rw-rw-r-- 1 ubuntu ubuntu 4.0G Jun 27 01:27 zero_pp_rank_0_mp_rank_00_optim_states.pt
-rw-rw-r-- 1 ubuntu ubuntu 4.0G Jun 27 01:27 zero_pp_rank_2_mp_rank_00_optim_states.pt
```
As you can see here file not present is "rank_5". I went back to earlier checkpoints and consistently `zero_pp_rank_5_mp_rank_00_optim_states.pt` is missing. I did monitor GPU usage during the training and nothing stuck me as odd. 

Let me report back with how the re-run goes.<|||||>First please file an issue with Deepspeed, since it's Deepspeed that saves these files. https://github.com/microsoft/DeepSpeed/issues

Are you running on all 8 gpus? can you validate in nvidia-smi? It should be the case since it reports:
```
Detected checkpoint of type zero stage 2, world_size: 8
```
Then I'd do debug logging (obviously switching to much more frequent save_steps, so it only takes one minute to log). I'd either:

* change this log into print so it's printed on each gpu
https://github.com/microsoft/DeepSpeed/blob/a029239812e15cf35334514449ed3127b915780a/deepspeed/runtime/engine.py#L1989

* or if you use `transformers` master and the HF trainer you can set ` --log_level info --log_level_replica info` and it'll log this info on each gpu w/o you needing to touch the deepspeed code.

For example with the above setting, on 2 gpus I get:
```
[2021-06-30 12:43:06,647] [INFO] [engine.py:1990:_save_zero_checkpoint] zero checkpoint saved output_dir/checkpoint-2/global_step2/zero_pp_rank_1_mp_rank_00_optim_states.pt
[2021-06-30 12:43:07,320] [INFO] [engine.py:1990:_save_zero_checkpoint] zero checkpoint saved output_dir/checkpoint-2/global_step2/zero_pp_rank_0_mp_rank_00_optim_states.pt
```

So we know both files got saved.<|||||>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/master/CONTRIBUTING.md) are likely to be ignored.