Repository: pytorch
Issue: Placing LSTM model on bfloat16 on GPU causes error
Labels: module: rnn,module: cuda,triaged,module: bfloat16

### 🐛 Describe the bug

import torch.nn as nn
import torch as th

# If using CPU as the device, the following codes run perfectly
rnn = nn.LSTM(10, 20, 2).to(device="cpu", dtype=th.bfloat16)
input = th.randn(5, 3, 10).to(device="cpu", dtype=th.bfloat16)
h0 = th.randn(2, 3, 20).to(device="cpu", dtype=th.bfloat16)
c0 = th.randn(2, 3, 20).to(device="cpu", dtype=th.bfloat16)
output, (hn, cn) = rnn(input, (h0, c0))

# However, change the device to be GPU, it gives out an error
rnn = nn.LSTM(10, 20, 2).to(device="cuda", dtype=th.bfloat16)
input = th.randn(5, 3, 10).to(device="cuda", dtype=th.bfloat16)
h0 = th.randn(2, 3, 20).to(device="cuda", dtype=th.bfloat16)
c0 = th.randn(2, 3, 20).to(device="cuda", dtype=th.bfloat16)
output, (hn, cn) = rnn(input, (h0, c0))
---------------------------------------------------------------------------
RuntimeError                              Traceback (most recent call last)
~\AppData\Local\Temp/ipykernel_20268/1164814484.py in <module>
      3 h0 = th.randn(2, 3, 20).to(device="cuda", dtype=th.bfloat16)
      4 c0 = th.randn(2, 3, 20).to(device="cuda", dtype=th.bfloat16)
----> 5 output, (hn, cn) = rnn(input, (h0, c0))

~\anaconda3\envs\py37\lib\site-packages\torch\nn\modules\module.py in _call_impl(self, *input, **kwargs)
   1108         if not (self._backward_hooks or self._forward_hooks or self._forward_pre_hooks or _global_backward_hooks
   1109                 or _global_forward_hooks or _global_forward_pre_hooks):
-> 1110             return forward_call(*input, **kwargs)
   1111         # Do not call functions when jit is used
   1112         full_backward_hooks, non_full_backward_hooks = [], []

~\anaconda3\envs\py37\lib\site-packages\torch\nn\modules\rnn.py in forward(self, input, hx)
    760         if batch_sizes is None:
    761             result = _VF.lstm(input, hx, self._flat_weights, self.bias, self.num_layers,
--> 762                               self.dropout, self.training, self.bidirectional, self.batch_first)
    763         else:
    764             result = _VF.lstm(input, batch_sizes, hx, self._flat_weights, self.bias,

RuntimeError: "_thnn_fused_lstm_cell_cuda" not implemented for 'BFloat16'

# I've also tried datatype 'float16', it works for GPU but doesn't work for CPU.
rnn = nn.LSTM(10, 20, 2).to(device="cpu", dtype=th.float16)
input = th.randn(5, 3, 10).to(device="cpu", dtype=th.float16)
h0 = th.randn(2, 3, 20).to(device="cpu", dtype=th.float16)
c0 = th.randn(2, 3, 20).to(device="cpu", dtype=th.float16)
output, (hn, cn) = rnn(input, (h0, c0))
---------------------------------------------------------------------------
RuntimeError                              Traceback (most recent call last)
~\AppData\Local\Temp/ipykernel_8944/1937453628.py in <module>
      3 h0 = th.randn(2, 3, 20).to(device="cpu", dtype=th.float16)
      4 c0 = th.randn(2, 3, 20).to(device="cpu", dtype=th.float16)
----> 5 output, (hn, cn) = rnn(input, (h0, c0))

~\anaconda3\envs\py37\lib\site-packages\torch\nn\modules\module.py in _call_impl(self, *input, **kwargs)
   1108         if not (self._backward_hooks or self._forward_hooks or self._forward_pre_hooks or _global_backward_hooks
   1109                 or _global_forward_hooks or _global_forward_pre_hooks):
-> 1110             return forward_call(*input, **kwargs)
   1111         # Do not call functions when jit is used
   1112         full_backward_hooks, non_full_backward_hooks = [], []

~\anaconda3\envs\py37\lib\site-packages\torch\nn\modules\rnn.py in forward(self, input, hx)
    760         if batch_sizes is None:
    761             result = _VF.lstm(input, hx, self._flat_weights, self.bias, self.num_layers,
--> 762                               self.dropout, self.training, self.bidirectional, self.batch_first)
    763         else:
    764             result = _VF.lstm(input, batch_sizes, hx, self._flat_weights, self.bias,

RuntimeError: "sigmoid_cpu" not implemented for 'Half'


### Versions

PyTorch version: 1.11.0
Is debug build: False
CUDA used to build PyTorch: 11.3
ROCM used to build PyTorch: N/A

OS: Microsoft Windows 11 Pro
GCC version: Could not collect
Clang version: Could not collect
CMake version: Could not collect
Libc version: N/A

Python version: 3.7.11 (default, Jul 27 2021, 09:42:29) [MSC v.1916 64 bit (AMD64)] (64-bit runtime)
Python platform: Windows-10-10.0.22000-SP0
Is CUDA available: True
CUDA runtime version: 11.7.64
GPU models and configuration: GPU 0: NVIDIA GeForce GTX 1650 Ti
Nvidia driver version: 516.59
cuDNN version: Could not collect
HIP runtime version: N/A
MIOpen runtime version: N/A

Versions of relevant libraries:
[pip3] mypy-extensions==0.4.3
[pip3] numpy==1.20.3
[pip3] numpy-ext==0.9.6
[pip3] torch==1.11.0
[pip3] torchaudio==0.11.0
[pip3] torchfile==0.1.0
[pip3] torchvision==0.12.0
[conda] blas                      1.0                         mkl    conda-forge
[conda] cudatoolkit               11.3.1               h59b6b97_2
[conda] mkl                       2021.4.0           haa95532_640
[conda] mkl-service               2.4.0            py37h2bbff1b_0
[conda] mkl_fft                   1.3.1            py37h277e83a_0
[conda] mkl_random                1.2.2            py37hf11a4ad_0
[conda] mypy-extensions           0.4.3                    pypi_0    pypi
[conda] numpy                     1.20.3           py37ha4e8547_0
[conda] numpy-base                1.20.3           py37hc2deb75_0
[conda] numpy-ext                 0.9.6                    pypi_0    pypi
[conda] pytorch                   1.11.0          py3.7_cuda11.3_cudnn8_0    pytorch
[conda] pytorch-mutex             1.0                        cuda    pytorch
[conda] torchaudio                0.11.0               py37_cu113    pytorch
[conda] torchfile                 0.1.0                      py_0    conda-forge
[conda] torchvision               0.12.0               py37_cu113    pytorch

cc @zou3519 @ngimel