pytorch autocast leads to baddbmm's output fp16 while input fp32.
Author: jovialchen
Pull request: no

### System Info

```
- `transformers` version: 4.29.1
- Platform: Windows-10-10.0.17763-SP0
- Python version: 3.9.16
- Huggingface_hub version: 0.14.1
- Safetensors version: 0.3.1
- PyTorch version (GPU?): 2.0.0+cu118 (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?: True
- Using distributed or parallel set-up in script?: False
```

### Who can help?

_No response_

### Information

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

### Tasks

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

### Reproduction

I'm trying to run 'bigcode/starcoder' and 'bigcode/starcoderbase-1b', but I'm facing an issue where the loss remains consistently zero. After debugging, I discovered that the inputs of the '`baddbmm`' function (specifically, '`attn_weights`', '`query`', and '`key`') are in fp32, while the outputs are in fp16. This discrepancy causes overflow in subsequent lines of code(assign `mask_value`(fp32) to `attn_weights`(fp16)), ultimately resulting in a loss value of 0.

During my investigation, I inserted a check using 'torch.is_autocast_enabled()' and found that it is actually set to TRUE in the code.

`dtype: (attn_weights=float32, query=float32,key=float32,beta=int,alpha=float), return attn_weights=float16
`
The code is located in the following file.
xxx\Anaconda3\envs\belle\Lib\site-packages\transformers\models\gpt_bigcode\modeling_gpt_bigcode.py

The scripts I'm using are https://github.com/hiyouga/LLaMA-Efficient-Tuning.

```
def _attn(self, query, key, value, attention_mask=None, head_mask=None):
        dtype = query.dtype
        softmax_dtype = torch.float32 if self.attention_softmax_in_fp32 else dtype
        upcast = dtype != softmax_dtype

        unscale = self.layer_idx + 1 if self.scale_attention_softmax_in_fp32 and upcast else 1
        scale_factor = unscale**-1
        if self.scale_attn_weights:
            scale_factor /= self.head_dim**0.5

        query_shape = query.shape
        batch_size = query_shape[0]
        key_length = key.size(-1)
        if self.multi_query:
            # (batch_size, query_length, num_heads, head_dim) x (batch_size, head_dim, key_length)
            # -> (batch_size, query_length, num_heads, key_length)
            query_length = query_shape[1]
            attn_shape = (batch_size, query_length, self.num_heads, key_length)
            attn_view = (batch_size, query_length * self.num_heads, key_length)
            # No copy needed for MQA 2, or when layer_past is provided.
            query = query.reshape(batch_size, query_length * self.num_heads, self.head_dim)
        else:
            # (batch_size, num_heads, query_length, head_dim) x (batch_size, num_heads, head_dim, key_length)
            # -> (batch_size, num_heads, query_length, key_length)
            query_length = query_shape[2]
            attn_shape = (batch_size, self.num_heads, query_length, key_length)
            attn_view = (batch_size * self.num_heads, query_length, key_length)
            # Always copies
            query = query.reshape(batch_size * self.num_heads, query_length, self.head_dim)
            # No copy when layer_past is provided.
            key = key.reshape(batch_size * self.num_heads, self.head_dim, key_length)

        attn_weights = torch.empty(attn_view, device=query.device, dtype=query.dtype)
        if query.device.type == "cpu":
            # This is needed because of a bug in pytorch https://github.com/pytorch/pytorch/issues/80588.
            # The bug was fixed in https://github.com/pytorch/pytorch/pull/96086,
            # but the fix has not been released as of pytorch version 2.0.0.
            attn_weights.zero_()
            beta = 1
        else:
            beta = 0
        **##!!error here!!**
        attn_weights = torch.baddbmm(attn_weights, query, key, beta=beta, alpha=scale_factor).view(attn_shape)
        ......
```



### Expected behavior

The input and output of '`baddbmm`' should be of the same dtype.

Comments:
I'll have a look, thanks for reporting. FYI @loubnabnl 

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.

Hey! Given that ` attn_weights = torch.empty(attn_view, device=query.device, dtype=query.dtype)` initialises the tensor based on the input type, and mostly given that we don't use autocast in our codebase I would suggest to disable auto_cast as a quick fix, and otherwise debug and check where the type casting comes from.
I don't really have the time to dig into your code and debug it for you so a minimal reproducer would also be welcomed ! 🤗  

Thank you, Arthur. I'll let you know if we have any findings.

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.