--save_on_each_node doesn't save pytorch_model.bin when deepspeed zero3 is used
Author: dunalduck0
Pull request: no
Labels: name: bug

### System Info

- `transformers` version: 4.20.1
- Platform: Linux-5.4.0-1083-azure-x86_64-with-glibc2.10
- Python version: 3.8.12
- Huggingface_hub version: 0.8.1
- PyTorch version (GPU?): 1.10.0a0+0aef44c (False)
- 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?: <fill in>
- Using distributed or parallel set-up in script?: <fill in>

I am finetuning GPT-J with DeepSpeed zero3 and 2 GPU nodes, each with 8 GPUs. I turned on --save_on_each_node but I found pytorch_model.bin is only saved by process world_rank=0 (i.e. process 0), not world_rank=8 (i.e. process 8). Below are the relevant part of logs from process 0 and process 8.

Process 0:
```
[INFO|trainer.py:2503] 2022-07-01 10:56:15,212 >> Saving model checkpoint to /tmp/azureml/cr/j/b531489bec744033bf0a29a63689aa44/cap/data-capability/wd/output_dir/checkpoint-10
[INFO|configuration_utils.py:446] 2022-07-01 10:56:15,213 >> Configuration saved in /tmp/azureml/cr/j/b531489bec744033bf0a29a63689aa44/cap/data-capability/wd/output_dir/checkpoint-10/config.json
[INFO|modeling_utils.py:1660] 2022-07-01 10:56:15,353 >> Model weights saved in /tmp/azureml/cr/j/b531489bec744033bf0a29a63689aa44/cap/data-capability/wd/output_dir/checkpoint-10/pytorch_model.bin
[INFO|tokenization_utils_base.py:2123] 2022-07-01 10:56:15,354 >> tokenizer config file saved in /tmp/azureml/cr/j/b531489bec744033bf0a29a63689aa44/cap/data-capability/wd/output_dir/checkpoint-10/tokenizer_config.json
[INFO|tokenization_utils_base.py:2130] 2022-07-01 10:56:15,354 >> Special tokens file saved in /tmp/azureml/cr/j/b531489bec744033bf0a29a63689aa44/cap/data-capability/wd/output_dir/checkpoint-10/special_tokens_map.json
[2022-07-01 10:56:20,868] [INFO] [engine.py:3229:save_16bit_model] Saving model weights to /tmp/azureml/cr/j/b531489bec744033bf0a29a63689aa44/cap/data-capability/wd/output_dir/checkpoint-10/pytorch_model.bin
[2022-07-01 10:56:43,320] [INFO] [logging.py:69:log_dist] [Rank 0] Saving model checkpoint: /tmp/azureml/cr/j/b531489bec744033bf0a29a63689aa44/cap/data-capability/wd/output_dir/checkpoint-10/global_step10/zero_pp_rank_0_mp_rank_00_model_states.pt
[2022-07-01 10:56:52,826] [INFO] [engine.py:3115:_save_zero_checkpoint] zero checkpoint saved /tmp/azureml/cr/j/b531489bec744033bf0a29a63689aa44/cap/data-capability/wd/output_dir/checkpoint-10/global_step10/zero_pp_rank_0_mp_rank_00_optim_states.pt
[INFO|trainer.py:1761] 2022-07-01 10:56:52,858 >> 
```

Process 8:
```
[INFO|trainer.py:2503] 2022-07-01 10:56:15,211 >> Saving model checkpoint to /tmp/azureml/cr/j/b531489bec744033bf0a29a63689aa44/cap/data-capability/wd/output_dir/checkpoint-10
[INFO|configuration_utils.py:446] 2022-07-01 10:56:15,213 >> Configuration saved in /tmp/azureml/cr/j/b531489bec744033bf0a29a63689aa44/cap/data-capability/wd/output_dir/checkpoint-10/config.json
[INFO|modeling_utils.py:1660] 2022-07-01 10:56:15,347 >> Model weights saved in /tmp/azureml/cr/j/b531489bec744033bf0a29a63689aa44/cap/data-capability/wd/output_dir/checkpoint-10/pytorch_model.bin
[INFO|tokenization_utils_base.py:2123] 2022-07-01 10:56:15,348 >> tokenizer config file saved in /tmp/azureml/cr/j/b531489bec744033bf0a29a63689aa44/cap/data-capability/wd/output_dir/checkpoint-10/tokenizer_config.json
[INFO|tokenization_utils_base.py:2130] 2022-07-01 10:56:15,348 >> Special tokens file saved in /tmp/azureml/cr/j/b531489bec744033bf0a29a63689aa44/cap/data-capability/wd/output_dir/checkpoint-10/special_tokens_map.json
[2022-07-01 10:56:52,831] [INFO] [engine.py:3115:_save_zero_checkpoint] zero checkpoint saved /tmp/azureml/cr/j/b531489bec744033bf0a29a63689aa44/cap/data-capability/wd/output_dir/checkpoint-10/global_step10/zero_pp_rank_8_mp_rank_00_optim_states.pt
[INFO|trainer.py:1761] 2022-07-01 10:56:52,858 >> 
```
The difference is that process 8 is missing these lines:
```
[2022-07-01 10:56:20,868] [INFO] [engine.py:3229:save_16bit_model] Saving model weights to /tmp/azureml/cr/j/b531489bec744033bf0a29a63689aa44/cap/data-capability/wd/output_dir/checkpoint-10/pytorch_model.bin
[2022-07-01 10:56:43,320] [INFO] [logging.py:69:log_dist] [Rank 0] Saving model checkpoint: /tmp/azureml/cr/j/b531489bec744033bf0a29a63689aa44/cap/data-capability/wd/output_dir/checkpoint-10/global_step10/zero_pp_rank_0_mp_rank_00_model_states.pt
```

My guess is that engine.py->save_16bit_model did not run on process 8. Looking at the [source code from DeepSpeed](https://github.com/microsoft/DeepSpeed/blob/9b70ce56e7af89d5226f9b06ebe1137407f371dc/deepspeed/runtime/engine.py#L3165), I suspect that these lines were not run because dist.get_rank() is not equal to 0 on process 8.
```
        if dist.get_rank() == 0:
            os.makedirs(save_dir, exist_ok=True)
            logger.info(f"Saving model weights to {path}")
            torch.save(state_dict, path)
```


### Who can help?

@stas00
@jeffra 

### 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

Not sure what to put here as the issue wouldn't happen unless more than one GPU node is used.

### Expected behavior

The main process on each node, or process 0 and process 8 in my examples, both save model file pytorch_model.bin

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