Wrong Data Type for Switch Transformer in Mixed Precision Training
Author: timlee0212
Pull request: no

### System Info

- `transformers` version: 4.32.1
- Platform: Linux-5.15.0-75-generic-x86_64-with-glibc2.35
- Python version: 3.10.13
- Huggingface_hub version: 0.17.3
- Safetensors version: 0.4.0
- Accelerate version: 0.24.1
- Accelerate config:    - compute_environment: LOCAL_MACHINE
        - distributed_type: MULTI_GPU
        - mixed_precision: bf16
        - use_cpu: False
        - debug: False
        - num_processes: 2
        - machine_rank: 0
        - num_machines: 1
        - gpu_ids: 0,1
        - rdzv_backend: static
        - same_network: True
        - main_training_function: main
        - downcast_bf16: no
        - tpu_use_cluster: False
        - tpu_use_sudo: False
        - tpu_env: []
- PyTorch version (GPU?): 2.1.0 (True)
- 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?: Yes
- Using distributed or parallel set-up in script?: Use DDP through accelerate and Trainer

### Who can help?

@ArthurZucker @younesbelkada 

### Information

- [X] The official example scripts
- [ ] My own modified scripts

### Tasks

- [X] An officially supported task in the `examples` folder (such as GLUE/SQuAD, ...)
- [ ] My own task or dataset (give details below)

### Reproduction

1. Run any training script with mixed-precision enabled (I enabled BF16 through `accelerate config` here)

    ```
    accelerate launch examples/pytorch/translation/run_translation.py
    --model_name_or_path google/switch-base-8     
    --do_train     
    --do_eval 
    --source_lang en     
    --target_lang de    
    --dataset_name wmt16     
    --dataset_config_name de-en    
     --output_dir /tmp/tst-translation     
    --per_device_train_batch_size=4     
    --per_device_eval_batch_size=4    
    --overwrite_output_dir     
    --predict_with_generate
    ```

2. Exception will happen,

```
 ...(Call Stack omitted)
  File "***/miniconda3/envs/moe/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1527, in _call_impl
    return forward_call(*args, **kwargs)
  File "***/miniconda3/envs/moe/lib/python3.10/site-packages/transformers/models/switch_transformers/modeling_switch_transformers.py", line 322, in forward
    next_states[token_indices] = expert(hidden_states[token_indices])
RuntimeError: Index put requires the source and destination dtypes match, got Float for the destination and BFloat16 for the source.
```

### Expected behavior

The training script should run without the dtype mismatch issue.

A workaround I found is:
In [src/transformers/models/switch_transformers/modeling_switch_transformers.py#L321](https://github.com/huggingface/transformers/blob/af3de8d87c717c4bb090f037d0d89413c195a42f/src/transformers/models/switch_transformers/modeling_switch_transformers.py#L321)

Change from:
```
            next_states[token_indices] = expert(hidden_states[token_indices])
```

to:

```
            next_states[token_indices] = expert(hidden_states[token_indices]).to(next_states.dtype)
```
The output of each expert is cast to BF16 (or FP16), but the hidden states are maintained in FP32. I think there should be one more step here to complete the conversion.

Comments:
Hi @timlee0212 
Thanks for the issue and the fix proposal, I think the fix you suggested makes sense, would you like to submit a PR for the fix? Otherwise, happy to do it!

Sure. I'm happy to submit a PR. 

Thanks for your in-time reply!