

TITLE
[wav2vec] deepspeed eval bug in the case of >1 gpus

COMMENTS
10

REACTIONS
+1: 1
-1: 0
laugh: 0
hooray: 0
heart: 0
rocket: 0
eyes: 0

BODY
## Environment info

- `transformers` version: 4.5.1
- Platform: Linux-4.15.0-140-generic-x86_64-with-debian-buster-sid
- Python version: 3.7.9
- PyTorch version (GPU?): 1.8.1 (True)
- Tensorflow version (GPU?): not installed (NA)
- Using GPU in script?: <2,4>
- Using distributed or parallel set-up in script?: <distributed>

### Who can help
@stas00 
@patrickvonplaten 
@patil-suraj 

## Information

I'm working on wav2vec2.0 using the following official script of huggingface.
https://github.com/huggingface/transformers/blob/master/examples/research_projects/wav2vec2/run_common_voice.py

I am trying to finetune huggingface model with multiple gpus using deepspeed.

```
deepspeed --num_gpus=1 run_common_voice.py --deepspeed ds_config.json --do_train --do_eval
```
works, but

```
deepspeed --num_gpus=2 run_common_voice.py --deepspeed ds_config.json --do_train --do_eval
```

stops working and freezes at the end of eval. 
The progress bar is 100% done but the eval result is not returned and it freezes.

## To reproduce

This is how to reproduce!
https://colab.research.google.com/drive/1VRCGcnhBlrMFYQ5aaNebucZuja-WB2I2?usp=sharing
Steps to reproduce the behavior:

1. Install deepspeed
2. Add `with autocast():` after line 481 in run_common_voice.py
3. Set param: `--deepspeed ds_config.json --do_train --do_eval`
4. Run run_common_voice.py using deepspeed with 1> gpus

ds_config has the following parameters.
```ds_config.json
{
  "fp16": {
    "enabled": "true",
    "loss_scale": 0,
    "loss_scale_window": 1000,
    "hysteresis": 2,
    "min_loss_scale": 1,
    "opt_level": "O3"
  },
  "steps_per_print": 100,
  "wall_clock_breakdown": "false"
}
```

## Expected behavior

The finetuning eval should be executed without freezing.

