Repository: pytorch
Issue: nn.linear not support bfloat16
Labels: needs reproduction,oncall: distributed,triaged

### 🐛 Describe the bug

from torch.distributed.fsdp import (
    FullyShardedDataParallel as FSDP,
    MixedPrecision,
    BackwardPrefetch,
    ShardingStrategy,
    FullStateDictConfig,
    StateDictType, CPUOffload,
)

class Model(nn.Module):
    def __init__(
        self,
    ):
        super().__init__()
    
        self.to_pred = nn.Linear(dim, 1, bias = False)


    def forward(
        self,
        input
    ):
       
        pred = self.to_pred(input)
        return pred

bfSixteen = MixedPrecision(
    param_dtype=torch.bfloat16,
    # Gradient communication precision.
    reduce_dtype=torch.bfloat16,
    # Buffer precision.
    buffer_dtype=torch.bfloat16,
)
model = Model()
model = FSDP(model,
                 auto_wrap_policy=opt_auto_wrap_policy,
                 mixed_precision=bfSixteen ,
     
                 device_id=torch.cuda.current_device())
res = model(input)




File "/nfs/users/weisongwei/miniconda3/envs/tf2.9/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1194, in _call_impl
    pred = self.to_pred(pooled)
      File "/nfs/users/weisongwei/miniconda3/envs/tf2.9/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1194, in _call_impl
pred = self.to_pred(pooled)
  File "/nfs/users/weisongwei/miniconda3/envs/tf2.9/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1194, in _call_impl
    return forward_call(*input, **kwargs)
  File "/nfs/users/weisongwei/miniconda3/envs/tf2.9/lib/python3.8/site-packages/torch/nn/modules/linear.py", line 114, in forward
    pred = self.to_pred(pooled)
  File "/nfs/users/weisongwei/miniconda3/envs/tf2.9/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1194, in _call_impl
    return forward_call(*input, **kwargs)
  File "/nfs/users/weisongwei/miniconda3/envs/tf2.9/lib/python3.8/site-packages/torch/nn/modules/linear.py", line 114, in forward
    return forward_call(*input, **kwargs)
  File "/nfs/users/weisongwei/miniconda3/envs/tf2.9/lib/python3.8/site-packages/torch/nn/modules/linear.py", line 114, in forward
        return forward_call(*input, **kwargs)return forward_call(*input, **kwargs)

  File "/nfs/users/weisongwei/miniconda3/envs/tf2.9/lib/python3.8/site-packages/torch/nn/modules/linear.py", line 114, in forward
  File "/nfs/users/weisongwei/miniconda3/envs/tf2.9/lib/python3.8/site-packages/torch/nn/modules/linear.py", line 114, in forward
    return forward_call(*input, **kwargs)
  File "/nfs/users/weisongwei/miniconda3/envs/tf2.9/lib/python3.8/site-packages/torch/nn/modules/linear.py", line 114, in forward
    return forward_call(*input, **kwargs)
  File "/nfs/users/weisongwei/miniconda3/envs/tf2.9/lib/python3.8/site-packages/torch/nn/modules/linear.py", line 114, in forward
    return forward_call(*input, **kwargs)
  File "/nfs/users/weisongwei/miniconda3/envs/tf2.9/lib/python3.8/site-packages/torch/nn/modules/linear.py", line 114, in forward
    return F.linear(input, self.weight, self.bias)
    RuntimeErrorreturn F.linear(input, self.weight, self.bias)    : 
    return F.linear(input, self.weight, self.bias)expected scalar type Float but found BFloat16return F.linear(input, self.weight, self.bias)

    
RuntimeErrorreturn F.linear(input, self.weight, self.bias)    : RuntimeError
return F.linear(input, self.weight, self.bias)RuntimeErrorexpected scalar type Float but found BFloat16: 
: RuntimeError
expected scalar type Float but found BFloat16expected scalar type Float but found BFloat16: 
RuntimeError
expected scalar type Float but found BFloat16    : 

### Versions

torch2.0 and torch1.13  

when use fsdp wrap a model which have nn.Linear , if mixed_precision use bfloat16, it will be an error.

cc @mrshenli @pritamdamania87 @zhaojuanmao @satgera @rohan-varma @gqchen @aazzolini @osalpekar @jiayisuse @H-Huang @kwen2501 @awgu @albanD @mruberry @jbschlosser @walterddr @mikaylagawarecki