Repository: pytorch
Issue: `torch.nn.LazyLinear` crash when using torch.bfloat16 dtype in pytorch 1.12.0 and 1.13.0
Labels: module: nn,triaged,intel

### 🐛 Describe the bug

In Pytorch 1.12.0 and 1.13.0, `torch.nn.LazyLinear` will crash when using `torch.bfloat16` dtype, but the same codes run well with `torch.float32`.

I think this issue may have a similar root cause to [#88658](https://github.com/pytorch/pytorch/issues/88658).

In addition, both `torch.nn.Linear` and `torch.nn.LazyLinear` have other interesting errors in Pytorch 1.10.0 and 1.11.0.
When they are using `torch.bfloat16`, they will crash with `RuntimeError: could not create a primitive`, which is a frequent error in PyTorch issues.
This also seems to be another unexpected behavior.

### To Reproduce
```python
from torch import nn

import torch
def test():
	tmp_result= torch.nn.LazyLinear(out_features=250880, bias=False, dtype=torch.bfloat16)
	return tmp_result
lm_head = test()
input=torch.ones(size=(8,1024,1536), dtype=torch.bfloat16)
output=lm_head(input)
```

The result in 1.12.0 and 1.13.0 is
```
/opt/conda/lib/python3.7/site-packages/torch/nn/modules/lazy.py:178: UserWarning: Lazy modules are a new feature under heavy development so changes to the API or functionality can happen at any moment.
  warnings.warn('Lazy modules are a new feature under heavy development '
Traceback (most recent call last):
  File "vanilla.py", line 9, in <module>
    output=lm_head(input)
  File "/opt/conda/lib/python3.7/site-packages/torch/nn/modules/module.py", line 1148, in _call_impl
    result = forward_call(*input, **kwargs)
  File "/opt/conda/lib/python3.7/site-packages/torch/nn/modules/linear.py", line 114, in forward
    return F.linear(input, self.weight, self.bias)
RuntimeError: [enforce fail at alloc_cpu.cpp:66] . DefaultCPUAllocator: can't allocate memory: you tried to allocate 263066747008 bytes. Error code 12 (Cannot allocate memory)
```

The result in 1.10.0 and 1.11.0 is
```
/opt/conda/lib/python3.7/site-packages/torch/nn/modules/lazy.py:178: UserWarning: Lazy modules are a new feature under heavy development so changes to the API or functionality can happen at any moment.
  warnings.warn('Lazy modules are a new feature under heavy development '
Traceback (most recent call last):
  File "vanilla.py", line 9, in <module>
    output=lm_head(input)
  File "/opt/conda/lib/python3.7/site-packages/torch/nn/modules/module.py", line 1120, in _call_impl
    result = forward_call(*input, **kwargs)
  File "/opt/conda/lib/python3.7/site-packages/torch/nn/modules/linear.py", line 103, in forward
    return F.linear(input, self.weight, self.bias)
  File "/opt/conda/lib/python3.7/site-packages/torch/nn/functional.py", line 1848, in linear
    return torch._C._nn.linear(input, weight, bias)
RuntimeError: could not create a primitive
```


### Expected behavior
Maybe it should have a similar performance with the codes with `float32` dtype, as shown in follows (It only raises a warning for Lazy module):

```
/opt/conda/lib/python3.7/site-packages/torch/nn/modules/lazy.py:178: UserWarning: Lazy modules are a new feature under heavy development so changes to the API or functionality can happen at any moment.
  warnings.warn('Lazy modules are a new feature under heavy development '
```

### Versions

I use docker to get Pytorch environments. 
`docker pull pytorch/pytorch:1.10.0-cuda11.3-cudnn8-devel;
docker pull pytorch/pytorch:1.12.0-cuda11.3-cudnn8-runtime;
docker pull pytorch/pytorch:1.13.0-cuda11.6-cudnn8-runtime` 
GPU models and configuration: RTX 3090

pytorch 1.10.0
```
[pip3] numpy==1.21.2
[pip3] torch==1.10.0
[pip3] torch-tb-profiler==0.4.0
[pip3] torchelastic==0.2.0
[pip3] torchtext==0.11.0
[pip3] torchvision==0.11.0
[conda] blas                      1.0                         mkl  
[conda] cudatoolkit               11.3.1               ha36c431_9    nvidia
[conda] ffmpeg                    4.3                  hf484d3e_0    pytorch
[conda] mkl                       2021.3.0           h06a4308_520  
[conda] mkl-service               2.4.0            py37h7f8727e_0  
[conda] mkl_fft                   1.3.1            py37hd3c417c_0  
[conda] mkl_random                1.2.2            py37h51133e4_0  
[conda] numpy                     1.21.2           py37h20f2e39_0  
[conda] numpy-base                1.21.2           py37h79a1101_0  
[conda] pytorch                   1.10.0          py3.7_cuda11.3_cudnn8.2.0_0    pytorch
[conda] pytorch-mutex             1.0                        cuda    pytorch
[conda] torch-tb-profiler         0.4.0                    pypi_0    pypi
[conda] torchelastic              0.2.0                    pypi_0    pypi
[conda] torchtext                 0.11.0                     py37    pytorch
[conda] torchvision               0.11.0               py37_cu113    pytorch
```


pytorch 1.12.0
```
[pip3] numpy==1.21.5
[pip3] torch==1.12.0
[pip3] torchelastic==0.2.0
[pip3] torchtext==0.13.0
[pip3] torchvision==0.13.0
[conda] blas                      1.0                         mkl  
[conda] cudatoolkit               11.3.1               ha36c431_9    nvidia
[conda] ffmpeg                    4.3                  hf484d3e_0    pytorch
[conda] mkl                       2021.4.0           h06a4308_640  
[conda] mkl-service               2.4.0            py37h7f8727e_0  
[conda] mkl_fft                   1.3.1            py37hd3c417c_0  
[conda] mkl_random                1.2.2            py37h51133e4_0  
[conda] numpy                     1.21.5           py37he7a7128_2  
[conda] numpy-base                1.21.5           py37hf524024_2  
[conda] pytorch                   1.12.0          py3.7_cuda11.3_cudnn8.3.2_0    pytorch
[conda] pytorch-mutex             1.0                        cuda    pytorch
[conda] torchelastic              0.2.0                    pypi_0    pypi
[conda] torchtext                 0.13.0                     py37    pytorch
[conda] torchvision               0.13.0               py37_cu113    pytorch
```

pytorch 1.13.0
```
[pip3] numpy==1.22.3
[pip3] torch==1.13.0
[pip3] torchtext==0.14.0
[pip3] torchvision==0.14.0
[conda] blas                      1.0                         mkl  
[conda] ffmpeg                    4.3                  hf484d3e_0    pytorch
[conda] mkl                       2021.4.0           h06a4308_640  
[conda] mkl-service               2.4.0            py39h7f8727e_0  
[conda] mkl_fft                   1.3.1            py39hd3c417c_0  
[conda] mkl_random                1.2.2            py39h51133e4_0  
[conda] numpy                     1.22.3           py39he7a7128_0  
[conda] numpy-base                1.22.3           py39hf524024_0  
[conda] pytorch                   1.13.0          py3.9_cuda11.6_cudnn8.3.2_0    pytorch
[conda] pytorch-cuda              11.6                 h867d48c_0    pytorch
[conda] pytorch-mutex             1.0                        cuda    pytorch
[conda] torchtext                 0.14.0                     py39    pytorch
[conda] torchvision               0.14.0               py39_cu116    pytorch
```

cc @albanD @mruberry @jbschlosser @walterddr @saketh-are @jgong5 @mingfeima @XiaobingSuper @sanchitintel @ashokei @jingxu10