[wav2vec] deepspeed eval bug in the case of >1 gpus
## 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.


deepspeed doesn't work with `autocast`, it has its own way of dealing with mixed precision, if you look in the `trainer.py` it's carefully bypassed. 

does the problem go away if you remove `autocast`?<|||||>@stas00  Thank you for your reply!
When I deleted `autocast` and ran it, I got the error `RuntimeError (Input type (torch.cuda.FloatTensor) and weight type (torch.cuda.HalfTensor) should be the same)`.
I ran it with autocast to eliminate this error.
FYI, when I do not do_eval or use only 1 GPU, the code run fine with autocast and deepspeed.

The full text of the error is below.

```
File "run_common_voice.py", line 512, in <module>
    main()
  File "run_common_voice.py", line 484, in main
    train_result = trainer.train(resume_from_checkpoint=checkpoint)
  File "/usr/local/lib/python3.7/dist-packages/transformers/trainer.py", line 1240, in train
    tr_loss += self.training_step(model, inputs)
  File "run_common_voice.py", line 232, in training_step
    loss = self.compute_loss(model, inputs)
  File "/usr/local/lib/python3.7/dist-packages/transformers/trainer.py", line 1667, in compute_loss
    outputs = model(**inputs)
  File "/usr/local/lib/python3.7/dist-packages/torch/nn/modules/module.py", line 889, in _call_impl
    result = self.forward(*input, **kwargs)
  File "/usr/local/lib/python3.7/dist-packages/deepspeed/runtime/engine.py", line 928, in forward
    loss = self.module(*inputs, **kwargs)
  File "/usr/local/lib/python3.7/dist-packages/torch/nn/modules/module.py", line 889, in _call_impl
    result = self.forward(*input, **kwargs)
  File "/usr/local/lib/python3.7/dist-packages/transformers/models/wav2vec2/modeling_wav2vec2.py", line 1050, in forward
    return_dict=return_dict,
  File "/usr/local/lib/python3.7/dist-packages/torch/nn/modules/module.py", line 889, in _call_impl
    result = self.forward(*input, **kwargs)
  File "/usr/local/lib/python3.7/dist-packages/transformers/models/wav2vec2/modeling_wav2vec2.py", line 828, in forward
    hidden_states = self.feature_extractor(input_values)
  File "/usr/local/lib/python3.7/dist-packages/torch/nn/modules/module.py", line 889, in _call_impl
    result = self.forward(*input, **kwargs)
  File "/usr/local/lib/python3.7/dist-packages/transformers/models/wav2vec2/modeling_wav2vec2.py", line 253, in forward
    hidden_states = conv_layer(hidden_states)
  File "/usr/local/lib/python3.7/dist-packages/torch/nn/modules/module.py", line 889, in _call_impl
    result = self.forward(*input, **kwargs)
  File "/usr/local/lib/python3.7/dist-packages/transformers/models/wav2vec2/modeling_wav2vec2.py", line 156, in forward
    hidden_states = self.conv(hidden_states)
  File "/usr/local/lib/python3.7/dist-packages/torch/nn/modules/module.py", line 889, in _call_impl
    result = self.forward(*input, **kwargs)
  File "/usr/local/lib/python3.7/dist-packages/torch/nn/modules/conv.py", line 263, in forward
    return self._conv_forward(input, self.weight, self.bias)
  File "/usr/local/lib/python3.7/dist-packages/torch/nn/modules/conv.py", line 260, in _conv_forward
    self.padding, self.dilation, self.groups)
RuntimeError: Input type (torch.cuda.FloatTensor) and weight type (torch.cuda.HalfTensor) should be the same
```

You can check this error in the last cell of the following colab.
https://colab.research.google.com/drive/1VRCGcnhBlrMFYQ5aaNebucZuja-WB2I2?usp=sharing<|||||>Thank you for making this reproducible, @tommy19970714 - I haven't worked yet with wav2vec, so I will have a look and get back to you.<|||||>OK, this is a new type of model that requires a special type of handling.

The NLP models get `long` inputs which get converted to the same dtype as the embedding weights, which under deepspeed/fp16 are `float16`. Currently deepspeed does `model.half`.

This model however receives inputs that are `float32` and it doesn't check whether the model weights are fp16 or not. Hence the error.

So this is one way to fix it:
```
diff --git a/src/transformers/models/wav2vec2/modeling_wav2vec2.py b/src/transformers/models/wav2vec2/modeling_wav2vec2.py
index 98123bdd3..639c2bc13 100755
--- a/src/transformers/models/wav2vec2/modeling_wav2vec2.py
+++ b/src/transformers/models/wav2vec2/modeling_wav2vec2.py
@@ -153,7 +153,7 @@ class Wav2Vec2LayerNormConvLayer(nn.Module):
         self.activation = ACT2FN[config.feat_extract_activation]

     def forward(self, hidden_states):
-        hidden_states = self.conv(hidden_states)
+        hidden_states = self.conv(hidden_states.to(dtype=self.conv.weight.dtype))

         hidden_states = hidden_states.transpose(-2, -1)
         hidden_states = self.layer_norm(hidden_states)
```

The test I was using is:
```
CUDA_VISIBLE_DEVICES=0 deepspeed --num_gpus=1 \
examples/research_projects/wav2vec2/run_common_voice.py \
--model_name_or_path="facebook/wav2vec2-large-xlsr-53" --dataset_config_name="tr" \
--output_dir=./wav2vec2-large-xlsr-turkish-demo --overwrite_output_dir --num_train_epochs="5" \
--per_device_train_batch_size="16" --learning_rate="3e-4" --warmup_steps="500" \
--evaluation_strategy="steps" --save_steps="5" --eval_steps="5" --logging_steps="5" \
--save_total_limit="3" --freeze_feature_extractor --feat_proj_dropout="0.0" --layerdrop="0.1" \
--gradient_checkpointing --fp16 --group_by_length --do_train --do_eval --deepspeed \
tests/deepspeed/ds_config_zero2.json
```

Could probably move it to the top-level layer so it'd work in all cases, if this exact path isn't always taken.

So this overcomes:
```
RuntimeError: Input type (torch.cuda.FloatTensor) and weight type (torch.cuda.HalfTensor) should be the same
```

but now running into:
```
  File "examples/research_projects/wav2vec2/run_common_voice.py", line 512, in <module>
    main()
  File "examples/research_projects/wav2vec2/run_common_voice.py", line 484, in main
    train_result = trainer.train(resume_from_checkpoint=checkpoint)
  File "/mnt/nvme1/code/huggingface/transformers-master/src/transformers/trainer.py", line 1240, in train
    tr_loss += self.training_step(model, inputs)
  File "examples/research_projects/wav2vec2/run_common_voice.py", line 232, in training_step
    loss = self.compute_loss(model, inputs)
  File "/mnt/nvme1/code/huggingface/transformers-master/src/transformers/trainer.py", line 1667, in compute_loss
    outputs = model(**inputs)
  File "/home/stas/anaconda3/envs/py38-pt18/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1015, in _call_impl
    return forward_call(*input, **kwargs)
  File "/mnt/nvme1/code/github/00optimize/deepspeed/deepspeed/runtime/engine.py", line 942, in forward
    loss = self.module(*inputs, **kwargs)
  File "/home/stas/anaconda3/envs/py38-pt18/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1015, in _call_impl
    return forward_call(*input, **kwargs)
  File "/mnt/nvme1/code/huggingface/transformers-master/src/transformers/models/wav2vec2/modeling_wav2vec2.py", line 1076, in forward
    loss = F.ctc_loss(
  File "/home/stas/anaconda3/envs/py38-pt18/lib/python3.8/site-packages/torch/nn/functional.py", line 2436, in ctc_loss
    return torch.ctc_loss(
RuntimeError: "ctc_loss_cuda" not implemented for 'Half'
```
so need to look more to see what to do there, probably need to switch to float32 just for that op.

However, it appears that may be this model can't be trained/eval'ed in fp16/mixed precision? 

When I run:
```
CUDA_VISIBLE_DEVICES=0 python examples/research_projects/wav2vec2/run_common_voice.py \
--model_name_or_path="facebook/wav2vec2-large-xlsr-53" --dataset_config_name="tr" \
--output_dir=./wav2vec2-large-xlsr-turkish-demo --overwrite_output_dir --num_train_epochs="5" \
--per_device_train_batch_size="16" --learning_rate="3e-4" --warmup_steps="500" \
--evaluation_strategy="steps" --save_steps="5" --eval_steps="5" --logging_steps="5" \
--save_total_limit="3" --freeze_feature_extractor --feat_proj_dropout="0.0" --layerdrop="0.1" \
--gradient_checkpointing --fp16 --group_by_length --do_train --do_eval
```
I see:
```
{'loss': nan, 'learning_rate': 4.2e-06, 'epoch': 0.05}  
```

We have multiple models that won't train under `fp16`-mixed precision, because they were pretrained in `bfloat16` which doesn't lend to `fp16` numerical range.

Deepspeed devs are working on adding the fp32 mode (next release hopefully). https://github.com/microsoft/DeepSpeed/pull/1004

p.s. please don't mix `amp` with running modes that don't use `amp` (deepspeed is one of them)  <|||||>Hi, @stas00 
Thanks for your help!
(I am working together with @tommy19970714 )

I saw your tweet about the new release of version 0.3.16.
https://github.com/microsoft/DeepSpeed/releases/tag/v0.3.16
https://huggingface.co/transformers/master/main_classes/trainer.html#fp32-precision

I set the `deepspeed.json` config to `auto`, referring to the article.

```JSON
{
  "fp16": {
    "enabled": "auto",
    "loss_scale": 0,
    "loss_scale_window": 1000,
    "hysteresis": 2,
    "min_loss_scale": 1,
    "opt_level": "O3"
  },
  "steps_per_print": 100,
  "wall_clock_breakdown": "false"
}
```

In addition to your suggestion, I made some changes to the model file to convert `log_probs` to float32:

```python
diff --git a/src/transformers/models/wav2vec2/modeling_wav2vec2.py b/src/transformers/models/wav2vec2/modeling_wav2vec2.py
index ba548dc3d..ce2ecdbe3 100755
--- a/src/transformers/models/wav2vec2/modeling_wav2vec2.py
+++ b/src/transformers/models/wav2vec2/modeling_wav2vec2.py
@@ -153,7 +153,7 @@ class Wav2Vec2LayerNormConvLayer(nn.Module):
         self.activation = ACT2FN[config.feat_extract_activation]

     def forward(self, hidden_states):
-        hidden_states = self.conv(hidden_states)
+        hidden_states = self.conv(hidden_states.to(dtype=self.conv.weight.dtype))

         hidden_states = hidden_states.transpose(-2, -1)
         hidden_states = self.layer_norm(hidden_states)
@@ -1071,10 +1071,15 @@ class Wav2Vec2ForCTC(Wav2Vec2PreTrainedModel):
             flattened_targets = labels.masked_select(labels_mask)

             log_probs = F.log_softmax(logits, dim=-1).transpose(0, 1)
+            # log_probs = log_probs.to(dtype=torch.float32), # doesn't work here

             with torch.backends.cudnn.flags(enabled=False):
                 loss = F.ctc_loss(
-                    log_probs,
+                    log_probs.to(dtype=torch.float32),
+                    # log_probs.to(dtype=torch.bfloat16),
+                    # log_probs,
                     flattened_targets,
                     input_lengths,
                     target_lengths,
```

Then it somehow worked! 
Does this seem to be a proper fix?
I might not be fully understanding the type differences tho.

Also, what else can I do to have it merged into the main branch?
I am willing to contribute, but I am not sure if the code is good enough. 
I assume it is missing config handling. <|||||>Thank you for suggested adjustments, @qqhann 

For the proper solution we shouldn't mess with the model ;)

The inputs `dtype` change is normally done inside the training loop, because it knows the context of the training. We just didn't need to do it until now, since as I mentioned earlier for NLP models we get the inputs adjusted to the right type through embedding lookup, so this is different.

One of the important parts here is to add tests for each of these situations.

What would be really useful is if you could help with creating a tiny wav2vec2 random model, to enable quick functional tests.

Here are some examples of such scripts:
- https://huggingface.co/stas/mt5-tiny-random/blob/main/mt5-make-tiny-model.py
- https://huggingface.co/stas/t5-very-small-random/blob/main/t5-make-very-small-model.py

In both cases it takes a normal model and reshapes it to a much smaller size. Usually the hard part is to figure out the non-model parts - dicts, tokenizers, etc. I don't know yet anything about wav2vec2 so it'd help if you had the know-how to create it.

The idea behind a tiny model is that it runs just like a normal model, but its weights are random, it's very small ~5-10MB or even smaller, it loads fast, and of course it produces random results. This is perfect for functional testing.

If you're not sure how to approach it, that's alright too. We will figure it out.<|||||>To update: @patrickvonplaten is kindly going to create a few tiny models and using his tiny `--dataset_name=patrickvonplaten/librispeech_asr_dummy` it should be possible to use `examples/research_projects/wav2vec2/run_asr.py` as the dev and test bench, so when this happens I should be able to complete this work. 

Until then your workaround is probably good enough if it's working for you.<|||||>You're welcome to follow my progress at fixing this issue at https://github.com/huggingface/transformers/pull/11638

ZeRO-2 works fully. ZeRO-3 still has one issue, but fp32 works. 

Do try and let me know if you run into any problems.

<|||||>@stas00 Thanks for letting me know! I'll keep an eye on it!<|||||>Update: with deepspeed master both zero-2 and zero-3 now work https://github.com/huggingface/transformers/pull/11638

It's ready to be merged.

Please give it a try.