training with native pytorch bf16 fsdp gives runtime error.
Author: YamamotoSayaka
Pull request: no

### System Info

- `transformers` version: 4.26.1
- Platform: Linux-3.10.0-1160.71.1.el7.x86_64-x86_64-with-glibc2.17
- Python version: 3.8.16
- Huggingface_hub version: 0.12.0
- PyTorch version (GPU?): 1.13.1 (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?: True

### Who can help?

@ArthurZucker 

### 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 train my bert model with native pytorch fsdp and bf16 mixed precision. 
When enabling bf16,  transformers would give me the follow error:

`File "/data/root/anaconda3/envs/pt/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1194, in _call_impl
    return forward_call(*input, **kwargs)
  File "/data/root/anaconda3/envs/pt/lib/python3.8/site-packages/transformers/models/bert/modeling_bert.py", line 363, in forward
    context_layer = torch.matmul(attention_probs, value_layer)
RuntimeError: expected scalar type BFloat16 but found Float
`

Below is my full py code
```
import os
import time
import functools
import numpy as np

import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim

import torchvision
from transformers import BertModel, BertConfig, BertLayer
from torchvision.models.vision_transformer import Encoder, EncoderBlock, MLPBlock

import torch.distributed as dist
import torch.multiprocessing as mp
from torch.nn.parallel import DistributedDataParallel as DDP
from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
from torch.distributed.fsdp import MixedPrecision

from torch.distributed.fsdp.wrap import (
    transformer_auto_wrap_policy,
    size_based_auto_wrap_policy,
    enable_wrap,
    wrap,
)
from torch.distributed.algorithms._checkpoint.checkpoint_wrapper import (
    apply_activation_checkpointing,
    checkpoint_wrapper as checkpoint_wrapper_pytorch,
    CheckpointImpl,
)


def setup(rank, world_size):
    os.environ['MASTER_ADDR'] = 'localhost'
    os.environ['MASTER_PORT'] = '12355'

    # initialize the process group
    dist.init_process_group("nccl", rank=rank, world_size=world_size)


def train(data, model, rank, world_size, optimizer):
    model.train()
    ddp_loss = torch.zeros(2).to(rank)
    data = data.to(rank)

    print("data: ", data.shape)

    for _ in range(8):
        optimizer.zero_grad()
        output = model(data)
        loss = output[0].mean()

        if rank == 0: print(loss)

        loss.backward()
        optimizer.step()
        ddp_loss[0] += loss.item()
        ddp_loss[1] += len(data)

    dist.all_reduce(ddp_loss, op=dist.ReduceOp.SUM)

def fsdp_main(rank, world_size):
    setup(rank, world_size)
    torch.cuda.set_device(rank)

    bs = 64
    model = BertModel(BertConfig())
    data = torch.ones((bs, 128)).int()

    attn_layers = [BertLayer]
    auto_wrap_policy = functools.partial(
        transformer_auto_wrap_policy,
        transformer_layer_cls = { *attn_layers }
    )
    dtype = torch.bfloat16
    mp_policy = MixedPrecision(
        param_dtype = dtype, reduce_dtype=dtype, buffer_dtype=dtype
    )
    # mp_policy = None
    model = FSDP(model.to(rank), device_id=rank,
            mixed_precision=mp_policy,
            auto_wrap_policy=auto_wrap_policy)
    # model = DDP(model.to(rank))
    if rank == 0:
        print(f"{model}")
    optimizer = optim.Adadelta(model.parameters(), lr=1e-2)
    train(data, model, rank, world_size, optimizer)
    dist.destroy_process_group()

if __name__ == '__main__':
    WORLD_SIZE = torch.cuda.device_count()
    # fsdp_main(0, WORLD_SIZE, args, )
    mp.spawn(fsdp_main,
        args=(WORLD_SIZE, ),
        nprocs=WORLD_SIZE,
        join=True)
```

You can launch the script with python run.py. Only pytorch and transformer are needed. I also found a similar report in the [pytorch repo](https://github.com/pytorch/pytorch/issues/75676). Looks like sth is wrong with transformers get_extended_mask implementation. Wondering if there will be any quick fix to this.

Many thanks!








### Expected behavior

Transformers should work with pytorch native bf16 

Comments:
I think the issue is just that you never do `model.half()` can you try this? 

Hi @ArthurZucker 
I tried running `model.half()` with `dtype = torch.float16`, and it works fine.
But when `dtype = 'bfloat16',` I still have to maunally convert the extended_mask output to bf16


Ok, as a minimal reproducing script here is what is working: 
```python 
from transformers import BertModel, BertTokenizer
import torch
model = BertModel.from_pretrained("bert-base-uncased")
tokenizer = BertTokenizer.from_pretrained("bert-base-uncased")
inputs = tokenizer("this is a test input", return_tensors = "pt")
model = model.cuda().to(torch.bfloat16)
print(model.dtype) # wether we can use it or not
model(**{k:v.cuda() for k,v in inputs.items()})
```
Also works with default config: 
```python 
model = BertModel(BertConfig())
model = model.cuda().to(torch.bfloat16)
print(model.dtype) # wether we can use it or not
model(**{k:v.cuda() for k,v in inputs.items()})
```

As I am not familiar with FSDP and it's internals, I need a full traceback error (because the `get_extended_attention_mask` uses `self.dtype` by default is no `dtype` is provided). Otherwise it might not be entirely related to transformers 😉 

cc @pacman100 for FSDP :-) 

Hello @YamamotoSayaka, does using context manager `torch.cuda.amp.autocast(dtype=torch.bfloat16)`   solve the issue? I tried it and seems to solve this issue. Just want a confirmation from your end.